diff --git a/flamedisx/__init__.py b/flamedisx/__init__.py index fa5169098..079f5b7c4 100644 --- a/flamedisx/__init__.py +++ b/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 diff --git a/flamedisx/sabre/__init__.py b/flamedisx/sabre/__init__.py new file mode 100644 index 000000000..30fff34b5 --- /dev/null +++ b/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 * diff --git a/flamedisx/sabre/sabre_blocks/__init__.py b/flamedisx/sabre/sabre_blocks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/flamedisx/sabre/sabre_blocks/energy_spectrum.py b/flamedisx/sabre/sabre_blocks/energy_spectrum.py new file mode 100644 index 000000000..8e79e4d44 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/flamedisx/sabre/sabre_blocks/final_signal.py b/flamedisx/sabre/sabre_blocks/final_signal.py new file mode 100644 index 000000000..bddc107d8 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/flamedisx/sabre/sabre_blocks/photon_production.py b/flamedisx/sabre/sabre_blocks/photon_production.py new file mode 100644 index 000000000..8717c84c3 --- /dev/null +++ b/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 diff --git a/flamedisx/sabre/sabre_models.py b/flamedisx/sabre/sabre_models.py new file mode 100644 index 000000000..412faa4aa --- /dev/null +++ b/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 = () \ No newline at end of file