Very rough skeleton for SABRE block model.

flamedisx/__init__.py

@@ -41,3 +41,7 @@
 # Custom TFP files
 # Access through fd.tfp_files.xxx
 from . import tfp_files
+
+# SABRE models
+# Access through fd.sabre.xxx
+from . import sabre

flamedisx/sabre/__init__.py

@@ -0,0 +1,5 @@
+from .sabre_blocks.energy_spectrum import *
+from .sabre_blocks.photon_production import *
+from .sabre_blocks.final_signal import *
+
+from .sabre_models import *

flamedisx/sabre/sabre_blocks/energy_spectrum.py

@@ -0,0 +1,95 @@
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+
+import os
+
+from multihist import Histdd
+import pickle as pkl
+
+import flamedisx as fd
+export, __all__ = fd.exporter()
+o = tf.newaxis
+
+
+class EnergySpectrum(fd.FirstBlock):
+    dimensions = ('energy',)
+    model_attributes = ('energies',)
+
+    #: Tensor listing energies this source can produce.
+    #: Approximate the energy spectrum as a sequence of delta functions.
+    energies = tf.cast(tf.linspace(1., 100., 100),
+                       dtype=fd.float_type())
+
+    def domain(self, data_tensor):
+        assert isinstance(self.energies, tf.Tensor)
+        return {self.dimensions[0]: tf.repeat(fd.np_to_tf(self.energies)[o, :],
+                                              self.source.batch_size,
+                                              axis=0)}
+
+    def _annotate(self, d):
+        pass
+
+    def random_truth(self, n_events, fix_truth=None, **params):
+        """Return pandas dataframe with event positions and times
+        randomly drawn.
+        """
+        data = pd.DataFrame()
+
+        spectrum_numpy = fd.tf_to_np(self.rates_vs_energy)
+        assert len(spectrum_numpy) == len(self.energies), \
+            "Energies and spectrum have different length"
+
+        data['energy'] = np.random.choice(
+            fd.tf_to_np(self.energies),
+            size=n_events,
+            p=spectrum_numpy / spectrum_numpy.sum(),
+            replace=True)
+        assert np.all(data['energy'] >= 0), "Generated negative energies??"
+
+        # For a constant-shape spectrum, fixing truth values is easy:
+        # we just overwrite the simulated values.
+        # Fixing one does not constrain any of the others.
+        self.source._overwrite_fixed_truths(data, fix_truth, n_events)
+
+        return data
+
+    def validate_fix_truth(self, d):
+        pass
+
+    def _compute(self, data_tensor, ptensor, **kwargs):
+        raise NotImplementedError
+
+    def mu_before_efficiencies(self, **params):
+        raise NotImplementedError
+
+
+@export
+class FixedShapeEnergySpectrum(EnergySpectrum):
+    """For a source whose energy spectrum has the same shape
+    throughout space and time.
+
+    If you add a rate variation with space, you must override draw_positions
+     and mu_before_efficiencies.
+     If you add a rate variation with time, you must override draw_times.
+     If you add a rate variation depending on both space and time, you must
+     override all of random_truth!
+
+    By default, this uses a flat 0 - 100 keV spectrum, sampled at 100 points.
+    """
+
+    model_attributes = ('rates_vs_energy',) + EnergySpectrum.model_attributes
+    model_functions = ()
+
+    #: Tensor listing the number of events for each energy the souce produces
+    #: Recall we approximate energy spectra by a sequence of delta functions.
+    rates_vs_energy = tf.ones(100, dtype=fd.float_type())
+
+    def _compute(self, data_tensor, ptensor, *, energy):
+        spectrum = tf.repeat(self.rates_vs_energy[o, :],
+                             self.source.batch_size,
+                             axis=0)
+        return spectrum
+
+    def mu_before_efficiencies(self, **params):
+        return np.sum(fd.np_to_tf(self.rates_vs_energy))

flamedisx/sabre/sabre_blocks/final_signal.py

@@ -0,0 +1,75 @@
+import typing as ty
+
+import numpy as np
+from scipy import stats
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+import flamedisx as fd
+export, __all__ = fd.exporter()
+o = tf.newaxis
+
+
+@export
+class MakeFinalSignal(fd.Block):
+    model_attributes = ()  # leave it explicitly empty
+
+    # Prevent pycharm warnings:
+    source: fd.Source
+    gimme: ty.Callable
+    gimme_numpy: ty.Callable
+
+    dimensions = ('integrated_charge', 'photons_produced')
+    special_model_functions = ()
+    model_functions = (
+        'photoelectron_gain_mean',
+        'photoelectron_gain_std',) + special_model_functions
+
+    max_dim_size = {'photons_produced': 120}
+
+    photoelectron_gain_mean = 1.
+    photoelectron_gain_std = 0.5
+
+    def _simulate(self, d):
+        d['integrated_charge'] = stats.norm.rvs(
+            loc=(d['photons_produced']
+                 * self.gimme_numpy('photoelectron_gain_mean')),
+            scale=(d['photons_produced']**0.5
+                   * self.gimme_numpy('photoelectron_gain_std')))
+
+    def _annotate(self, d):
+        m = self.gimme_numpy('photoelectron_gain_mean')
+        s = self.gimme_numpy('photoelectron_gain_std')
+
+        mle = d['photons_produced_mle'] = \
+            (d['integrated_charge'] / m).clip(0, None)
+        scale = mle**0.5 * s / m
+
+        for bound, sign, intify in (('min', -1, np.floor),
+                                    ('max', +1, np.ceil)):
+            # For detected quanta the MLE is quite accurate
+            # (since fluctuations are tiny)
+            # so let's just use the relative error on the MLE)
+            d['photons_produced_' + bound] = intify(
+                mle + sign * self.source.max_sigma * scale
+            ).clip(0, None).astype(int)
+
+    def _compute(self,
+                 data_tensor, ptensor,
+                 photons_produced, integrated_charge):
+        # Lookup signal gain mean and std per detected quanta
+        mean_per_q = self.gimme('photoelectron_gain_mean',
+                                data_tensor=data_tensor,
+                                ptensor=ptensor)[:, o, o]
+        std_per_q = self.gimme('photoelectron_gain_std',
+                               data_tensor=data_tensor,
+                               ptensor=ptensor)[:, o, o]
+
+        mean = photons_produced * mean_per_q
+        std = photons_produced ** 0.5 * std_per_q
+
+        # add offset to std to avoid NaNs from norm.pdf if std = 0
+        result = tfp.distributions.Normal(
+            loc=mean, scale=std + 1e-10
+        ).prob(integrated_charge)
+        return result

flamedisx/sabre/sabre_blocks/photon_production.py

@@ -0,0 +1,42 @@
+import numpy as np
+from scipy import stats
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+import flamedisx as fd
+export, __all__ = fd.exporter()
+o = tf.newaxis
+
+
+@export
+class MakePhotons(fd.Block):
+
+    depends_on = ((('energy',), 'rate_vs_energy'),)
+    dimensions = ('photons_produced', 'energy')
+
+    special_model_functions = ('light_yield',)
+    model_functions = special_model_functions
+
+    light_yield = 1.   # Nonsense, source provides specifics
+
+    def _compute(self,
+                 data_tensor, ptensor,
+                 # Domain
+                 photons_produced,
+                 # Dependency domain and value
+                 energy, rate_vs_energy):
+        ly = self.source.gimme('light_yield', bonus_arg=energy,
+                                data_tensor=data_tensor, ptensor=ptensor)
+        mean_yield = energy * ly
+
+        return rate_vs_energy[:, o, :] * tfp.distributions.Poisson(
+                rate=mean_yield).prob(photons_produced)
+
+    def _simulate(self, d):
+        d['ly'] = self.gimme_numpy('light_yield', bonus_arg=d['energy'].values)
+
+        d['mean_yield'] = d['ly'] * d['energy']
+        d['photons_produced'] = stats.poisson.rvs(mu=d['mean_yield'])
+
+    def _annotate(self, d):
+        pass

flamedisx/sabre/sabre_models.py

@@ -0,0 +1,24 @@
+import tensorflow as tf
+
+import flamedisx as fd
+from .. import sabre as fd_sabre
+
+export, __all__ = fd.exporter()
+
+
+@export
+class SABRESource(fd.BlockModelSource):
+    model_blocks = (
+        fd_sabre.FixedShapeEnergySpectrum,
+        fd_sabre.MakePhotons,
+        fd_sabre.MakeFinalSignal)
+
+    @staticmethod
+    def light_yield(energy, *, abs_ly=10.):
+        """
+        """
+        ly = abs_ly
+        return ly
+
+    final_dimensions = ('integrated_charge',)
+    no_step_dimensions = ()