Fixed adding multiple constraint terms to the log likelihood

FlamTeam · Jan 29, 2024 · c87b4b0 · c87b4b0
1 parent d50b217
commit c87b4b0
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py
@@ -523,10 +523,15 @@ def _log_likelihood(self,
             0.)
         if dsetname == self.dsetnames[0]:
             if constraint_extra_args is None:
-                ll += self.log_constraint(**params_unstacked)
+                ll += tf.where( tf.equal(i_batch, tf.constant(0, dtype=fd.int_type())),
+                                self.log_constraint(**params_unstacked),
+                                0.)
             else:
                 kwargs = {**params_unstacked, **constraint_extra_args}
-                ll += self.log_constraint(**kwargs)
+                ll += tf.where( tf.equal(i_batch, tf.constant(0, dtype=fd.int_type())),
+                                self.log_constraint(**kwargs),
+                                0.)
+
 
         # Autodifferentiation. This is why we use tensorflow:
         grad = tf.gradients(ll, grad_par_stack)[0]