add acceleration option to JointPrimaryMarginalizedModel likelihood

bin/inference/pycbc_inference

@@ -141,7 +141,12 @@ with ctx:
     if pool.is_main_process():
         for fn in [sampler.checkpoint_file, sampler.backup_file]:
             with loadfile(fn, 'a') as fp:
-                fp.write_config_file(cp)
+                # some models will interally modify original cp for sampling,
+                # such as joint_primary_marginalized, we need to save original
+                if hasattr(model, 'original_config'):
+                    fp.write_config_file(model.original_config)
+                else:
+                    fp.write_config_file(cp)
 
     # Run the sampler
     sampler.run()

bin/inference/pycbc_inference_model_stats

@@ -96,7 +96,10 @@ model.sampling_transforms = None
 def callmodel(arg):
     iteration, paramvals = arg
     # calculate the logposterior to get all stats populated
-    model.update(**{p: paramvals[p] for p in model.variable_params})
+    try:
+        model.update_all_models(**{p: paramvals[p] for p in model.variable_params})
+    except:
+        model.update(**{p: paramvals[p] for p in model.variable_params})
     _ = model.logposterior
     stats = model.get_current_stats()
 

pycbc/inference/models/hierarchical.py

@@ -607,39 +607,33 @@ def _loglikelihood(self):
 
 
 class JointPrimaryMarginalizedModel(HierarchicalModel):
-    """ Hierarchical heterodyne likelihood for coherent multiband
-    parameter estimation which combines data from space-borne and
-    ground-based GW detectors coherently. Currently, this only
-    supports LISA as the space-borne GW detector.
-
-    Sub models are treated as if the same GW source (such as a GW
-    from stellar-mass BBH) is observed in different frequency bands by
-    space-borne and ground-based GW detectors, then transform all
-    the parameters into the same frame in the sub model level, use
-    `HierarchicalModel` to get the joint likelihood, and marginalize
-    over all the extrinsic parameters supported by `RelativeTimeDom`
-    or its variants. Note that LISA submodel only supports the `Relative`
-    for now, for ground-based detectors, please use `RelativeTimeDom`
-    or its variants.
-
-    Although this likelihood model is used for multiband parameter
-    estimation, users can still use it for other purposes, such as
-    GW + EM parameter estimation, in this case, please use `RelativeTimeDom`
-    or its variants for the GW data, for the likelihood of EM data,
-    there is no restrictions.
+    """This likelihood model can be used for cases when one of the submodels
+    can be marginalized to accelerate the total likelihood. This likelihood
+    model also allows for further acceleration of other models during
+    marginalization, if some extrinsic parameters can be tightly constrained
+    by the primary model. More specifically, such as the EM + GW parameter
+    estimation, the sky localization can be well measured. For LISA + 3G
+    multiband observation, SOBHB signals' (tc, ra, dec) can be tightly
+    constrained by 3G network, so this model is also useful for this case.
     """
     name = 'joint_primary_marginalized'
 
     def __init__(self, variable_params, submodels, **kwargs):
         super().__init__(variable_params, submodels, **kwargs)
 
+        # store the original config to self
+        self.original_config = kwargs['original_config']
         # assume the ground-based submodel as the primary model
         self.primary_model = self.submodels[kwargs['primary_lbl'][0]]
         self.primary_lbl = kwargs['primary_lbl'][0]
         self.other_models = self.submodels.copy()
         self.other_models.pop(kwargs['primary_lbl'][0])
         self.other_models = list(self.other_models.values())
 
+        # determine whether to accelerate total_loglr
+        self.static_margin_params_in_other_models = \
+            'static_margin_params_in_other_models' in kwargs
+
     def write_metadata(self, fp, group=None):
         """Adds metadata to the output files
 
@@ -652,6 +646,9 @@ def write_metadata(self, fp, group=None):
             by group, i.e., to ``fp[group].attrs``. Otherwise, metadata is
             written to the top-level attrs (``fp.attrs``).
         """
+        # replace the internal config for top-level model with
+        # the original config
+        fp.write_config_file(self.original_config)
         super().write_metadata(fp, group=group)
         sampattrs = fp.getattrs(group=fp.samples_group)
         # if a group is specified, prepend the lognl names with it
@@ -686,25 +683,48 @@ def total_loglr(self):
         """
         # calculate <d-h|d-h> = <h|h> - 2<h|d> + <d|d> up to a constant
 
-        # note that for SOBHB signals, ground-based detectors dominant SNR
-        # and accuracy of (tc, ra, dec)
         self.primary_model.return_sh_hh = True
         sh_primary, hh_primary = self.primary_model.loglr
         self.primary_model.return_sh_hh = False
+        # set logr, otherwise it will store (sh, hh)
+        setattr(self.primary_model._current_stats, 'loglr',
+                self.primary_model.marginalize_loglr(sh_primary, hh_primary))
 
         margin_names_vector = list(
             self.primary_model.marginalize_vector_params.keys())
         if 'logw_partial' in margin_names_vector:
             margin_names_vector.remove('logw_partial')
-
         margin_params = {}
-        nums = 1
-        for key, value in self.primary_model.current_params.items():
-            # add marginalize_vector_params
-            if key in margin_names_vector:
-                margin_params[key] = value
-                if isinstance(value, numpy.ndarray):
-                    nums = len(value)
+
+        if self.static_margin_params_in_other_models:
+            # Due to the high precision of extrinsic parameters constrined
+            # by the primary model, the mismatch of wavefroms in others by
+            # varing those parameters is pretty small, so we can keep them
+            # static to accelerate total_loglr. Here, we use matched-filering
+            # SNR instead of lilkelihood, because luminosity distance and
+            # inclination has a very strong degeneracy, change of inclination
+            # will change best match distance, so change the amplitude of
+            # waveform. Using SNR will cancel out the effect of amplitude.err
+            i_max_extrinsic = numpy.argmax(
+                numpy.abs(sh_primary) / hh_primary**0.5)
+            for p in margin_names_vector:
+                if isinstance(self.primary_model.current_params[p],
+                              numpy.ndarray):
+                    margin_params[p] = \
+                        self.primary_model.current_params[p][i_max_extrinsic]
+                    nums = len(self.primary_model.current_params[p])
+                else:
+                    margin_params[p] = self.primary_model.current_params[p]
+                    nums = 1
+        else:
+            for key, value in self.primary_model.current_params.items():
+                # add marginalize_vector_params
+                if key in margin_names_vector:
+                    margin_params[key] = value
+                    if isinstance(value, numpy.ndarray):
+                        nums = len(value)
+                    else:
+                        nums = 1
         # add distance if it has been marginalized,
         # use numpy array for it is just let it has the same
         # shape as marginalize_vector_params, here we assume
@@ -713,7 +733,7 @@ def total_loglr(self):
             margin_params['distance'] = numpy.full(
                 nums, self.primary_model.current_params['distance'])
 
-        # add likelihood contribution from space-borne detectors, we
+        # add likelihood contribution from other_models, we
         # calculate sh/hh for each marginalized parameter point
         sh_others = numpy.full(nums, 0 + 0.0j)
         hh_others = numpy.zeros(nums)
@@ -723,24 +743,47 @@ def total_loglr(self):
             # not using self.primary_model.current_params, because others_model
             # may have its own static parameters
             current_params_other = other_model.current_params.copy()
-            for i in range(nums):
+            if not self.static_margin_params_in_other_models:
+                for i in range(nums):
+                    current_params_other.update(
+                        {key: value[i] if isinstance(value, numpy.ndarray)
+                         else value for key, value in margin_params.items()})
+                    other_model.update(**current_params_other)
+                    other_model.return_sh_hh = True
+                    sh_other, hh_other = other_model.loglr
+                    sh_others[i] += sh_other
+                    hh_others[i] += hh_other
+                    other_model.return_sh_hh = False
+                    # set logr, otherwise it will store (sh, hh)
+                    setattr(other_model._current_stats, 'loglr',
+                            other_model.marginalize_loglr(sh_other, hh_other))
+            else:
+                # use one margin point set to approximate all the others
                 current_params_other.update(
-                    {key: value[i] if isinstance(value, numpy.ndarray) else
-                        value for key, value in margin_params.items()})
+                    {key: value[0] if isinstance(value, numpy.ndarray)
+                     else value for key, value in margin_params.items()})
                 other_model.update(**current_params_other)
                 other_model.return_sh_hh = True
-                sh_others[i], hh_others[i] = other_model.loglr
+                sh_other, hh_other = other_model.loglr
                 other_model.return_sh_hh = False
+                # set logr, otherwise it will store (sh, hh)
+                setattr(other_model._current_stats, 'loglr',
+                        other_model.marginalize_loglr(sh_other, hh_other))
+                sh_others += sh_other
+                hh_others += hh_other
 
         if nums == 1:
+            # the type of the original sh/hh_others are numpy.array,
+            # might not the same as sh/hh_primary during reconstruct,
+            # during reconstruct of distance, sh/hh_others need to be scalar
             sh_others = sh_others[0]
+            hh_others = hh_others[0]
         sh_total = sh_primary + sh_others
         hh_total = hh_primary + hh_others
 
-        # calculate marginalize_vector_weights
-        self.primary_model.marginalize_vector_weights = \
-            - numpy.log(self.primary_model.vsamples)
         loglr = self.primary_model.marginalize_loglr(sh_total, hh_total)
+        setattr(self._current_stats, 'total_loglr', loglr)
+
         return loglr
 
     def others_lognl(self):
@@ -795,6 +838,9 @@ def from_config(cls, cp, **kwargs):
         # we need the read from config function from the init; to prevent
         # circular imports, we import it here
         from pycbc.inference.models import read_from_config
+        # store the original config file, here use deeocopy to avoid later
+        # changes of cp affect it
+        kwargs['original_config'] = cp.__deepcopy__(cp)
         # get the submodels
         kwargs['primary_lbl'] = shlex.split(cp.get('model', 'primary_model'))
         kwargs['others_lbls'] = shlex.split(cp.get('model', 'other_models'))
@@ -805,6 +851,10 @@ def from_config(cls, cp, **kwargs):
         sparam_map = map_params(hpiter(cp.options('static_params'),
                                        submodel_lbls))
 
+        # get the acceleration label
+        kwargs['static_margin_params_in_other_models'] = shlex.split(
+            cp.get('model', 'static_margin_params_in_other_models'))
+
         # we'll need any waveform transforms for the initializing sub-models,
         # as the underlying models will receive the output of those transforms
 
@@ -856,18 +906,21 @@ def from_config(cls, cp, **kwargs):
                           cp.get('static_params', param.fullname))
 
             # set the variable params: different from the standard
-            # hierarchical model, in this multiband model, all sub-models
-            # has the same variable parameters, so we don't need to worry
-            # about the unique variable issue. Besides, the primary model
-            # needs to do marginalization, so we must set variable_params
-            # and prior section before initializing it.
+            # hierarchical model, in this JointPrimaryMarginalizedModel model,
+            # all sub-models has the same variable parameters, so we don't
+            # need to worry about the unique variable issue. Besides,
+            # the primary model needs to do marginalization, so we must set
+            # variable_params and prior section before initializing it.
 
             subcp.add_section('variable_params')
             for param in vparam_map[lbl]:
                 if lbl in kwargs['primary_lbl']:
+                    # set variable_params for the primary model
                     subcp.set('variable_params', param.subname,
                               cp.get('variable_params', param.fullname))
                 else:
+                    # all variable_params in other models will come
+                    # from the primary model during sampling
                     subcp.set('static_params', param.subname, 'REPLACE')
 
             for section in cp.sections():
@@ -919,14 +972,13 @@ def from_config(cls, cp, **kwargs):
         # it will not be listed in `variable_params` and `prior` sections
         primary_model = submodels[kwargs['primary_lbl'][0]]
         marginalized_params = primary_model.marginalize_vector_params.copy()
-        if 'logw_partial' in marginalized_params:
-            marginalized_params.pop('logw_partial')
-            marginalized_params = list(marginalized_params.keys())
-        else:
-            marginalized_params = []
-        # this may also include 'f_ref', 'f_lower', 'approximant',
-        # but doesn't matter
-        marginalized_params += list(primary_model.static_params.keys())
+        marginalized_params = list(marginalized_params.keys())
+        # add distance or phase if they are marginalized
+        if primary_model.distance_marginalization:
+            marginalized_params.append('distance')
+        if primary_model.marginalize_phase:
+            marginalized_params.append('coa_phase')
+
         for p in primary_model.static_params.keys():
             p_full = '%s__%s' % (kwargs['primary_lbl'][0], p)
             if p_full not in cp['static_params']:
@@ -940,6 +992,10 @@ def from_config(cls, cp, **kwargs):
                     cp['variable_params'].pop(p)
                     cp.pop(section)
 
+        # save the vitual config file to disk for later check
+        with open('internal_top.ini', 'w', encoding='utf-8') as file:
+            cp.write(file)
+
         # now load the model
         logging.info("Loading joint_primary_marginalized model")
         return super(HierarchicalModel, cls).from_config(

pycbc/inference/models/relbin.py

@@ -596,10 +596,11 @@ def _loglr(self):
             filt += filter_i
             norm += norm_i
 
+        loglr = self.marginalize_loglr(filt, norm)
         if self.return_sh_hh:
             results = (filt, norm)
         else:
-            results = self.marginalize_loglr(filt, norm)
+            results = loglr
         return results
 
     def write_metadata(self, fp, group=None):