WIP: An attempt at allowing the 'TCB' units in PINT

pint/models/astrometry.py

@@ -42,7 +42,6 @@ def __init__(self):
     def setup(self):
         super(Astrometry, self).setup()
 
-
     def ssb_to_psb_xyz_ICRS(self, epoch=None):
         """Returns unit vector(s) from SSB to pulsar system barycenter under ICRS.
 
@@ -190,6 +189,7 @@ def setup(self):
                             "POSEPOCH or PEPOCH are required if PM is set.")
                 else:
                     self.POSEPOCH.quantity = self.PEPOCH.quantity
+            self.POSEPOCH.scale = self.UNITS.value.lower()
 
     def print_par(self):
         result = ''
@@ -358,6 +358,7 @@ def setup(self):
                             "POSEPOCH or PEPOCH are required if PM is set.")
                 else:
                     self.POSEPOCH.quantity = self.PEPOCH.quantity
+                    self.POSEPOCH.scale = self.UNITS.value.lower()
 
     def get_psr_coords(self, epoch=None):
         """Returns pulsar sky coordinates as an astropy ecliptic oordinates

pint/models/binary_ell1.py

@@ -61,10 +61,13 @@ def setup(self):
                 warn("Since ECC is 0.0, using T0 as TASC.")
                 if self.T0.value is not None:
                     self.TASC.value = self.T0.value
+                    self.TASC.scale = self.UNITS.value.lower()
                 else:
                     raise MissingParameter("ELL1", 'T0', "T0 or TASC is required for ELL1 model.")
             else:
                 raise MissingParameter("ELL1", 'TASC', "TASC is required for ELL1 model.")
+        else:
+            self.TASC.scale = self.UNITS.value.lower()
 
 
 

pint/models/dispersion_model.py

@@ -80,6 +80,8 @@ def setup(self):
             if self.DMEPOCH.value is None:
                 raise MissingParameter("Dispersion", "DMEPOCH",
                         "DMEPOCH is required if DM1 or higher are set")
+            else:
+                self.DMEPOCH.scale = self.UNITS.value.lower()
         base_dms = list(self.get_prefix_mapping_component('DM').values())
         base_dms += ['DM',]
 
@@ -105,10 +107,10 @@ def base_dm(self, toas):
         dm = np.zeros(len(tbl))
         dm_terms = self.get_DM_terms()
         if self.DMEPOCH.value is None:
-            DMEPOCH = tbl['tdbld'][0]
+            DMEPOCH = tbl[self.UNITS.value.lower()+'ld'][0]
         else:
             DMEPOCH = self.DMEPOCH.value
-        dt = (tbl['tdbld'] - DMEPOCH) * u.day
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - DMEPOCH) * u.day
         dt_value = (dt.to(u.yr)).value
         dm_terms_value = [d.value for d in dm_terms]
         dm = taylor_horner(dt_value, dm_terms_value)
@@ -156,10 +158,10 @@ def d_delay_d_DMs(self, toas, param_name, acc_delay=None): # NOTE we should have
         dm_terms = np.longdouble(np.zeros(len(dms)))
         dm_terms[order] = np.longdouble(1.0)
         if self.DMEPOCH.value is None:
-            DMEPOCH = tbl['tdbld'][0]
+            DMEPOCH = tbl[self.UNITS.value.lower()+'ld'][0]
         else:
             DMEPOCH = self.DMEPOCH.value
-        dt = (tbl['tdbld'] - DMEPOCH) * u.day
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - DMEPOCH) * u.day
         dt_value = (dt.to(u.yr)).value
         d_dm_d_dm_param = taylor_horner(dt_value, dm_terms)* (self.DM.units/par.units)
         return DMconst * d_dm_d_dm_param/ bfreq**2.0

pint/models/glitch.py

@@ -76,6 +76,7 @@ def setup(self):
             if not hasattr(self, 'GLEP_%d' % idx):
                 msg = 'Glicth Epoch is needed for Glicth %d.' % idx
                 raise MissingParameter("Glitch", 'GLEP_%d' % idx, msg)
+            getattr(self, 'GLEP_%d' % idx).scale = self.UNITS.value.lower()
             for param in self.glitch_prop:
                 if not hasattr(self, param + '%d' % idx):
                     param0 = getattr(self, param + '1')
@@ -121,7 +122,7 @@ def glitch_phase(self, toas, delay):
                 dF0 = getattr(self, "GLF0_%d" % idx).quantity
                 dF1 = getattr(self, "GLF1_%d" % idx).quantity
                 dF2 = getattr(self, "GLF2_%d" % idx).quantity
-                dt = (tbl['tdbld'] - eph) * u.day - delay
+                dt = (tbl[self.UNITS.value.lower()+'ld'] - eph) * u.day - delay
                 dt = dt.to(u.second)
                 affected = dt > 0.0  # TOAs affected by glitch
                 # decay term
@@ -146,7 +147,7 @@ def d_phase_d_GLPH(self, toas, param, delay):
             raise ValueError("Can not calculate d_phase_d_GLPH with respect to %s." % param)
         eph = time_to_longdouble(getattr(self, "GLEP_" + ids).value)
         par_GLPH = getattr(self, param)
-        dt = (tbl['tdbld'] - eph) * u.day - delay
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - eph) * u.day - delay
         dt = dt.to(u.second)
         affected = numpy.where(dt > 0.0)[0]
         dpdGLPH = numpy.zeros(len(tbl), dtype=numpy.longdouble) * u.cycle/par_GLPH.units
@@ -163,7 +164,7 @@ def d_phase_d_GLF0(self, toas, param, delay):
             raise ValueError("Can not calculate d_phase_d_GLF0 with respect to %s." % param)
         eph = time_to_longdouble(getattr(self, "GLEP_" + ids).value)
         par_GLF0 = getattr(self, param)
-        dt = (tbl['tdbld'] - eph) * u.day - delay
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - eph) * u.day - delay
         dt = dt.to(u.second)
         affected = numpy.where(dt > 0.0)[0]
         dpdGLF0 = numpy.zeros(len(tbl), dtype=numpy.longdouble) * u.cycle/par_GLF0.units
@@ -179,7 +180,7 @@ def d_phase_d_GLF1(self, toas,  param, delay):
             raise ValueError("Can not calculate d_phase_d_GLF1 with respect to %s." % param)
         eph = time_to_longdouble(getattr(self, "GLEP_" + ids).value)
         par_GLF1 = getattr(self, param)
-        dt = (tbl['tdbld'] - eph) * u.day - delay
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - eph) * u.day - delay
         dt = dt.to(u.second)
         affected = numpy.where(dt > 0.0)[0]
         dpdGLF1 = numpy.zeros(len(tbl), dtype=numpy.longdouble) * u.cycle/par_GLF1.units
@@ -195,7 +196,7 @@ def d_phase_d_GLF2(self, toas,  param, delay):
             raise ValueError("Can not calculate d_phase_d_GLF2 with respect to %s." % param)
         eph = time_to_longdouble(getattr(self, "GLEP_" + ids).value)
         par_GLF2 = getattr(self, param)
-        dt = (tbl['tdbld'] - eph) * u.day - delay
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - eph) * u.day - delay
         dt = dt.to(u.second)
         affected = numpy.where(dt > 0.0)[0]
         dpdGLF2 = numpy.zeros(len(tbl), dtype=numpy.longdouble) * u.cycle/par_GLF2.units
@@ -213,7 +214,7 @@ def d_phase_d_GLF0D(self, toas, param, delay):
         eph = time_to_longdouble(getattr(self, "GLEP_" + ids).value)
         par_GLF0D = getattr(self, param)
         tau = getattr(self, "GLTD_%d" % idv).quantity
-        dt = (tbl['tdbld'] - eph) * u.day - delay
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - eph) * u.day - delay
         dt = dt.to(u.second)
         affected = numpy.where(dt > 0.0)[0]
         dpdGLF0D = numpy.zeros(len(tbl), dtype=numpy.longdouble) * u.cycle/par_GLF0D.units
@@ -235,7 +236,7 @@ def d_phase_d_GLTD(self, toas, param, delay):
             return numpy.zeros(len(tbl), dtype=numpy.longdouble) * u.cycle/par_GLTD.units
         glf0d = getattr(self, 'GLF0D_'+ids).quantity
         tau = par_GLTD.quantity
-        dt = (tbl['tdbld'] - eph) * u.day - delay
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - eph) * u.day - delay
         dt = dt.to(u.second)
         affected = numpy.where(dt > 0.0)[0]
         dpdGLTD = numpy.zeros(len(tbl), dtype=numpy.longdouble) * u.cycle/par_GLTD.units

pint/models/noise_model.py

@@ -181,7 +181,7 @@ def ecorr_basis_weight_pair(self, toas):
 
         """
         tbl = toas.table
-        t = (tbl['tdbld'].quantity * u.day).to(u.s).value
+        t = (tbl[self.UNITS.value.lower()+'ld'].quantity * u.day).to(u.s).value
         ecorrs = self.get_ecorrs()
         Umats = []
         for ec in ecorrs:
@@ -263,7 +263,7 @@ def pl_rn_basis_weight_pair(self, toas):
 
         """
         tbl = toas.table
-        t = (tbl['tdbld'].quantity * u.day).to(u.s).value
+        t = (tbl[self.UNITS.value.lower()+'ld'].quantity * u.day).to(u.s).value
         amp, gam, nf = self.get_pl_vals()
         Fmat, f = create_fourier_design_matrix(t, nf)
         weight = powerlaw(f, amp, gam) * f[0]

pint/models/pulsar_binary.py

@@ -32,7 +32,6 @@ def __init__(self,):
             units=u.day,
             description="Orbital period", long_double=True))
 
-
         self.add_param(p.floatParameter(name="PBDOT",
             units = u.day/u.day,
             description="Orbital period derivitve respect to time", \
@@ -89,6 +88,11 @@ def __init__(self,):
 
     def setup(self):
         super(PulsarBinary, self).setup()
+
+        if self.T0.value is not None:
+            # set the units
+            self.T0.scale = self.UNITS.value.lower()
+
         for bpar in self.params:
             self.register_deriv_funcs(self.d_binary_delay_d_xxxx, bpar)
         # Setup the model isinstance
@@ -137,17 +141,17 @@ def update_binary_object(self, toas, acc_delay=None):
         """
         # Don't need to fill P0 and P1. Translate all the others to the format
         # that is used in bmodel.py
-        # Get barycnetric toa first
+        # Get barycentric toa first
         updates = {}
         tbl = toas.table
         if acc_delay is None:
             # If the accumulate delay is not provided, it will try to get
             # the barycentric correction.
             acc_delay = self.delay(toas, self.__class__.__name__, False)
-        self.barycentric_time = tbl['tdbld'] * u.day - acc_delay
+        self.barycentric_time = tbl[self.UNITS.value.lower()+'ld'] * u.day - acc_delay
+        updates['psr_pos'] = self.ssb_to_psb_xyz_ICRS(epoch=tbl[self.UNITS.value.lower()+'ld'].astype(np.float64))
         updates['barycentric_toa'] = self.barycentric_time
         updates['obs_pos'] = tbl['ssb_obs_pos'].quantity
-        updates['psr_pos'] = self.ssb_to_psb_xyz_ICRS(epoch=tbl['tdbld'].astype(np.float64))
         for par in self.binary_instance.binary_params:
             binary_par_names = [par,]
             if par in self.binary_instance.param_aliases.keys():
@@ -183,7 +187,7 @@ def binarymodel_delay(self, toas, acc_delay=None):
         return self.binary_instance.binary_delay()
 
     def d_binary_delay_d_xxxx(self, toas, param, acc_delay):
-        """Return the bianry model delay derivtives"""
+        """Return the binary model delay derivtives"""
         self.update_binary_object(toas, acc_delay)
         return self.binary_instance.d_binarydelay_d_par(param)
 

pint/models/solar_system_shapiro.py

@@ -78,7 +78,7 @@ def solar_system_shapiro_delay(self, toas, acc_delay=None):
             if key['obs'].lower() == 'barycenter':
                 log.debug("Skipping Shapiro delay for Barycentric TOAs")
                 continue
-            psr_dir = self.ssb_to_psb_xyz_ICRS(epoch=grp['tdbld'].astype(numpy.float64))
+            psr_dir = self.ssb_to_psb_xyz_ICRS(epoch=grp[self.UNITS.value.lower()+'ld'].astype(numpy.float64))
             delay[loind:hiind] += self.ss_obj_shapiro_delay(grp['obs_sun_pos'],
                                     psr_dir, self._ss_mass_sec['sun'])
             if self.PLANET_SHAPIRO.value:

pint/models/solar_wind_dispersion.py

@@ -42,7 +42,7 @@ def solar_wind_delay(self, toas, acc_delay=None):
                 bfreq = tbl['freq']
 
             rsa = tbl['obs_sun_pos'].quantity
-            pos = self.ssb_to_psb_xyz_ICRS(epoch=tbl['tdbld'].astype(np.float64))
+            pos = self.ssb_to_psb_xyz_ICRS(epoch=tbl[self.UNITS.value.lower()+'ld'].astype(np.float64))
             r = np.sqrt(np.sum(rsa*rsa, axis=1))
             cos_theta = np.sum(rsa*pos, axis=1) / r
             ret =  const.au**2.0 * np.arccos(cos_theta) * DMconst * self.NE_SW.quantity / \
@@ -63,7 +63,7 @@ def d_delay_d_ne_sw(self, toas, param_name, acc_delay=None):
                 bfreq = tbl['freq']
 
             rsa = tbl['obs_sun_pos'].quantity
-            pos = self.ssb_to_psb_xyz_ICRS(epoch=tbl['tdbld'].astype(np.float64))
+            pos = self.ssb_to_psb_xyz_ICRS(epoch=tbl[self.UNITS.value.lower()+'ld'].astype(np.float64))
             r = np.sqrt(np.sum(rsa*rsa, axis=1))
             cos_theta = np.sum(rsa*pos, axis=1) / r
 

pint/models/spindown.py

@@ -44,6 +44,10 @@ def setup(self):
             if getattr(self, p).value is None:
                 raise MissingParameter("Spindown", p)
 
+        if self.PEPOCH.value is not None:
+            # set the units
+            self.PEPOCH.scale = self.UNITS.value.lower()
+
         # Check continuity
         F_terms = list(self.get_prefix_mapping_component('F').keys())
         F_terms.sort()
@@ -88,11 +92,11 @@ def get_dt(self, toas, delay):
         """
         tbl = toas.table
         if self.PEPOCH.value is None:
-            phsepoch_ld = time_to_longdouble(tbl['tdb'][0] - delay[0])
+            phsepoch_ld = time_to_longdouble(tbl[self.UNITS.value.lower()][0] - delay[0])
         else:
             phsepoch_ld = time_to_longdouble(self.PEPOCH.quantity)
 
-        dt = (tbl['tdbld'] - phsepoch_ld)*u.day - delay
+        dt = (tbl[self.UNITS.value.lower()+'ld'] - phsepoch_ld)*u.day - delay
 
         return dt
 

pint/models/timing_model.py

@@ -63,8 +63,8 @@ class TimingModel(object):
         A list for register the timing model component type name. For example,
         delay components will be register as 'DelayComponent'.
     top_level_params : list
-        A parameter name list for thoes parameters belong to the top timing
-        model class rather than a specific component.
+        A parameter name list for those parameters that belong to the top
+        timing model class rather than a specific component.
 
     Properties
     ----------
@@ -96,7 +96,7 @@ def __init__(self, name='', components=[]):
         self.add_param_from_top(strParameter(name="EPHEM",
             description="Ephemeris to use"), '')
         self.add_param_from_top(strParameter(name="UNITS",
-            description="Units (TDB assumed)"), '')
+            description="Units", value="TDB"), '')  # defaults to TDB if not given
 
         self.setup_components(components)
 
@@ -378,7 +378,7 @@ def add_param_from_top(self, param, target_component):
         else:
             if target_component not in list(self.components.keys()):
                 raise AttributeError("Can not find component '%s' in "
-                                     "timging model." % target_component)
+                                     "timing model." % target_component)
             self.components[target_component].add_param(param)
 
     def remove_param(self, param):
@@ -424,6 +424,7 @@ def get_params_of_type(self, param_type):
                 param_type.upper() == par_prefix.upper():
                 result.append(par.name)
         return result
+
     def get_prefix_mapping(self, prefix):
         """Get the index mapping for the prefix parameters.
            Parameter
@@ -564,8 +565,6 @@ def noise_model_basis_weight(self, toas):
             result.append(nf(toas)[1])
         return np.hstack((r for r in result))
 
-
-
     def get_barycentric_toas(self, toas, cutoff_component=''):
         """This is a convenient function for calculate the barycentric TOAs.
            Parameter
@@ -587,7 +586,12 @@ def get_barycentric_toas(self, toas, cutoff_component=''):
                 if cp.category == 'pulsar_system':
                     cutoff_component = cp.__class__.__name__
         corr = self.delay(toas, cutoff_component, False)
-        return tbl['tdbld'] * u.day - corr
+
+        timescale = 'TDB' if not hasattr(self, 'UNITS') else getattr(self, 'UNITS').value
+        if timescale in ['TDB', 'TCB']:
+            return tbl[timescale.lower()+'ld'] * u.day - corr
+        else:
+            raise ValueError("Unknown 'UNITS' value")
 
     def d_phase_d_toa(self, toas, sample_step=None):
         """Return the derivative of phase wrt TOA
@@ -791,12 +795,18 @@ def read_parfile(self, filename):
 
             k = l.split()
             name = k[0].upper()
-
-            if name == 'UNITS' and len(k) > 1 and k[1] != 'TDB':
-                log.error("UNITS %s not yet supported by PINT" % k[1])
-                raise Exception("UNITS %s not yet supported by PINT" % k[1])
 
-
+            if name == 'UNITS' and len(k) > 1:
+                if k[1] != 'TDB' and k[1] != 'TCB' and k[1] != 'SI':
+                    log.error("UNITS {} not yet supported by "
+                              "PINT".format(k[1]))
+                    raise Exception("UNITS {} not yet supported by "
+                                    "PINT".format(k[1]))
+
+                if k[1] == 'SI':
+                    # SI is just an alias for TCB so set to 'TCB'
+                    k[1] = 'TCB'
+
             if name in checked_param:
                 if name in repeat_param.keys():
                     repeat_param[name] += 1
@@ -1042,6 +1052,10 @@ def is_in_parfile(self,para_dict):
         """
         pNames_inpar = list(para_dict.keys())
         pNames_inModel = self.params
+
+        # remove UNITS as many components contain it
+        if 'UNITS' in pNames_inModel:
+            del pNames_inModel[pNames_inModel.index('UNITS')]
 
         # For solar system Shapiro delay component
         if hasattr(self,'PLANET_SHAPIRO'):
@@ -1085,6 +1099,7 @@ def print_par(self,):
             result += getattr(self, p).as_parfile_line()
         return result
 
+
 class DelayComponent(Component):
     def __init__(self,):
         super(DelayComponent, self).__init__()

pint/models/wave.py

@@ -33,6 +33,7 @@ def setup(self):
                                        "WAVE_OM is set.")
             else:
                 self.WAVEEPOCH = self.PEPOCH
+            self.WAVEEPOCH.scale = self.UNITS.value.lower()
 
         wave_terms = list(self.get_prefix_mapping_component('WAVE').keys())
         wave_terms.sort()
@@ -62,7 +63,8 @@ def wave_delay(self, toas, acc_delay=None):
                       range(1, self.num_wave_terms + 1)]
         wave_terms = [getattr(self, name) for name in wave_names]
         wave_om = self.WAVE_OM.quantity
-        time = self.barycentric_time = toas.table['tdbld'] * u.day
+        time = self.barycentric_time = toas.table[self.UNITS.value.lower()+'ld'] * u.day
+
         for k, wave_term in enumerate(wave_terms):
             wave_a, wave_b = wave_term.quantity
             k = k + 1