diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py index 10979dad6..15f0c4204 100644 --- a/flamedisx/likelihood.py +++ b/flamedisx/likelihood.py @@ -284,7 +284,7 @@ def set_data(self, continue # Copy ensures annotations don't clobber - source.set_data(deepcopy(data[dname]), data_is_annotated) + source.set_data(deepcopy(data[dname]), data_is_annotated,ignore_priors=True) # Update batch info dset_index = self.dsetnames.index(dname) @@ -355,6 +355,10 @@ def simulate(self, fix_truth=None, **params): mu = rm * s.mu_before_efficiencies( **self._filter_source_kwargs(params, sname)) # Simulate this many events from source + if type(self.sources[sname]).__name__!='TemplateSource': + #only do for non-template sources + mu=mu/self.mu_estimators[sname](**self._filter_source_kwargs(params, sname)) + n_to_sim = np.random.poisson(mu) if n_to_sim == 0: continue @@ -532,7 +536,6 @@ def _log_likelihood(self, self.log_constraint(**kwargs), 0.) - # Autodifferentiation. This is why we use tensorflow: grad = tf.gradients(ll, grad_par_stack)[0] if second_order: diff --git a/flamedisx/nest/lxe_blocks/energy_spectrum.py b/flamedisx/nest/lxe_blocks/energy_spectrum.py index 4de8462de..480883394 100644 --- a/flamedisx/nest/lxe_blocks/energy_spectrum.py +++ b/flamedisx/nest/lxe_blocks/energy_spectrum.py @@ -14,7 +14,7 @@ class EnergySpectrum(fd.FirstBlock): dimensions = ('energy',) model_attributes = ( - 'energies', + 'energies','max_dim_size', 'radius', 'z_top', 'z_bottom', 'z_topDrift', 'drift_velocity', 't_start', 't_stop', @@ -32,6 +32,14 @@ class EnergySpectrum(fd.FirstBlock): # Just a dummy 0-10 keV spectrum energies = tf.cast(tf.linspace(0., 10., 1000), dtype=fd.float_type()) + default_size=100 + max_dim_size : dict + def setup(self): + """ + Setups up a managable max dim size + """ + assert isinstance(self.max_dim_size,dict) + self.max_dim_size={'energy':self.max_dim_size['energy']} drift_map_dt = None def derive_drift_time(self, data): @@ -58,7 +66,7 @@ def derive_observed_xy(self, data): y_obs = self.drift_map_x(np.array([data['r'], data['z']]).T) * sinTheta return x_obs, y_obs - + def domain(self, data_tensor): assert isinstance(self.energies, tf.Tensor) # see WIMPsource for why @@ -283,15 +291,16 @@ def _compute(self, data_tensor, ptensor, *, energy): def mu_before_efficiencies(self, **params): return np.sum(fd.np_to_tf(self.rates_vs_energy)) - +#Relics from before dimsize was an attribute-maybe this can be changed? @export class FixedShapeEnergySpectrumNR(FixedShapeEnergySpectrum): - max_dim_size = {'energy': 150} - - + def setup(self): + self.max_dim_size=self.max_dim_size @export class FixedShapeEnergySpectrumER(FixedShapeEnergySpectrum): - max_dim_size = {'energy': 100} + def setup(self): + self.max_dim_size=self.max_dim_size + @export diff --git a/flamedisx/nest/lxe_blocks/quanta_splitting.py b/flamedisx/nest/lxe_blocks/quanta_splitting.py index 3ff53651a..02f797779 100644 --- a/flamedisx/nest/lxe_blocks/quanta_splitting.py +++ b/flamedisx/nest/lxe_blocks/quanta_splitting.py @@ -47,7 +47,6 @@ def _compute(self, ions_produced, # Dependency domain and value energy, rate_vs_energy): - def compute_single_energy(args, approx=False): # Compute the block for a single energy. # Set approx to True for an approximate computation at higher energies @@ -62,7 +61,12 @@ def compute_single_energy(args, approx=False): # Calculate the ion domain tensor for this energy _ions_produced = ions_produced_add + ions_min - + #every event in the batch shares E therefore ions domain + _ions_produced_1D=_ions_produced[0,0,0,:] + #create p_ni/p_nq domain for ER/NR + nq_2D=tf.repeat(unique_quanta[:,o],tf.shape(_ions_produced_1D)[0],axis=1) + ni_2D=tf.repeat(_ions_produced_1D[o,:],tf.shape(unique_quanta)[0],axis=0) + if self.is_ER: nel_mean = self.gimme('mean_yield_electron', data_tensor=data_tensor, ptensor=ptensor, bonus_arg=energy) @@ -72,19 +76,25 @@ def compute_single_energy(args, approx=False): bonus_arg=nq_mean) if approx: - p_nq = tfp.distributions.Normal(loc=nq_mean, - scale=tf.sqrt(nq_mean * fano) + 1e-10).prob(nq) + p_nq_1D = tfp.distributions.Normal(loc=nq_mean, + scale=tf.sqrt(nq_mean * fano) + 1e-10).prob(unique_quanta) else: normal_dist_nq = tfp.distributions.Normal(loc=nq_mean, - scale=tf.sqrt(nq_mean * fano) + 1e-10) - p_nq = normal_dist_nq.cdf(nq + 0.5) - normal_dist_nq.cdf(nq - 0.5) + scale=tf.sqrt(nq_mean * fano) + 1e-10) + p_nq_1D=normal_dist_nq.cdf(unique_quanta + 0.5) - normal_dist_nq.cdf(unique_quanta - 0.5) + #restore p_ni from unique_nq x n_ions -> unique_nel x n_electrons x n_photons (does not need n_ions) + p_nq=tf.gather_nd(params=p_nq_1D,indices=index_nq[:,o],batch_dims=0) + p_nq=tf.reshape(p_nq,[tf.shape(nq)[0],tf.shape(nq)[1],tf.shape(nq)[2]]) + ex_ratio = self.gimme('exciton_ratio', data_tensor=data_tensor, ptensor=ptensor, bonus_arg=energy) alpha = 1. / (1. + ex_ratio) - p_ni = tfp.distributions.Binomial( - total_count=nq, probs=alpha).prob(_ions_produced) + p_ni_2D=tfp.distributions.Binomial(total_count=nq_2D, probs=alpha).prob(ni_2D) + #restore p_ni from unique_nq x n_ions -> unique_nel x n_electrons x n_photons x n_ions + p_ni=tf.gather_nd(params=p_ni_2D,indices=index_nq[:,o],batch_dims=0) + p_ni=tf.reshape(tf.reshape(p_ni,[-1]),[tf.shape(nq)[0],tf.shape(nq)[1],tf.shape(nq)[2],tf.shape(nq)[3]]) else: yields = self.gimme('mean_yields', data_tensor=data_tensor, ptensor=ptensor, @@ -98,37 +108,47 @@ def compute_single_energy(args, approx=False): bonus_arg=nq_mean) ni_fano = yield_fano[0] nex_fano = yield_fano[1] - + + #p_ni does not need to have an altered dimensionality, see tensordot. if approx: - p_ni = tfp.distributions.Normal(loc=nq_mean*alpha, - scale=tf.sqrt(nq_mean*alpha*ni_fano) + 1e-10).prob(_ions_produced) + p_ni_1D = tfp.distributions.Normal(loc=nq_mean*alpha, + scale=tf.sqrt(nq_mean*alpha*ni_fano) + 1e-10).prob(_ions_produced_1D) - p_nq = tfp.distributions.Normal(loc=nq_mean*alpha*ex_ratio, + p_nq_2D = tfp.distributions.Normal(loc=nq_mean*alpha*ex_ratio, scale=tf.sqrt(nq_mean*alpha*ex_ratio*nex_fano) + 1e-10).prob( - nq - _ions_produced) + nq_2D - ni_2D) else: normal_dist_ni = tfp.distributions.Normal(loc=nq_mean*alpha, scale=tf.sqrt(nq_mean*alpha*ni_fano) + 1e-10) - p_ni = normal_dist_ni.cdf(_ions_produced + 0.5) - \ - normal_dist_ni.cdf(_ions_produced - 0.5) + p_ni_1D = normal_dist_ni.cdf(_ions_produced_1D + 0.5) - \ + normal_dist_ni.cdf(_ions_produced_1D - 0.5) normal_dist_nq = tfp.distributions.Normal(loc=nq_mean*alpha*ex_ratio, scale=tf.sqrt(nq_mean*alpha*ex_ratio*nex_fano) + 1e-10) - p_nq = normal_dist_nq.cdf(nq - _ions_produced + 0.5) \ - - normal_dist_nq.cdf(nq - _ions_produced - 0.5) + p_nq_2D = normal_dist_nq.cdf(nq_2D - ni_2D + 0.5) \ + - normal_dist_nq.cdf(nq_2D - ni_2D - 0.5) + + # restore p_nq from unique_nq x n_ions -> unique_nel x n_electrons x n_photons x n_ions + p_nq=tf.gather_nd(params=p_nq_2D,indices=index_nq[:,o],batch_dims=0) + p_nq=tf.reshape(tf.reshape(p_nq,[-1]),[tf.shape(nq)[0],tf.shape(nq)[1],tf.shape(nq)[2],tf.shape(nq)[3]]) + + + + nel_2D=tf.repeat(unique_nel[:,o],tf.shape(_ions_produced_1D)[0],axis=1) + ni_nel_2D=tf.repeat(_ions_produced_1D[o,:],tf.shape(unique_nel)[0],axis=0) recomb_p = self.gimme('recomb_prob', data_tensor=data_tensor, ptensor=ptensor, bonus_arg=(nel_mean, nq_mean, ex_ratio)) skew = self.gimme('skewness', data_tensor=data_tensor, ptensor=ptensor, bonus_arg=nq_mean) var = self.gimme('variance', data_tensor=data_tensor, ptensor=ptensor, - bonus_arg=(nel_mean, nq_mean, recomb_p, _ions_produced)) + bonus_arg=(nel_mean, nq_mean, recomb_p, ni_nel_2D)) width_corr = self.gimme('width_correction', data_tensor=data_tensor, ptensor=ptensor, bonus_arg=skew) mu_corr = self.gimme('mu_correction', data_tensor=data_tensor, ptensor=ptensor, bonus_arg=(skew, var, width_corr)) - mean = (tf.ones_like(_ions_produced, dtype=fd.float_type()) - recomb_p) * _ions_produced - mu_corr + mean = (tf.ones_like(ni_nel_2D, dtype=fd.float_type()) - recomb_p) * ni_nel_2D - mu_corr std_dev = tf.sqrt(var) / width_corr if self.is_ER: @@ -137,19 +157,28 @@ def compute_single_energy(args, approx=False): owens_t_terms = 5 if approx: - p_nel = fd.tfp_files.SkewGaussian(loc=mean, scale=std_dev, - skewness=skew, - owens_t_terms=owens_t_terms).prob(electrons_produced) + p_nel_1D = fd.tfp_files.SkewGaussian(loc=mean, scale=std_dev, + skewness=skew, + owens_t_terms=owens_t_terms).prob(nel_2D) else: - p_nel = fd.tfp_files.TruncatedSkewGaussianCC(loc=mean, scale=std_dev, - skewness=skew, - limit=_ions_produced, - owens_t_terms=owens_t_terms).prob(electrons_produced) - - p_mult = p_nq * p_ni * p_nel - - # Contract over ions_produced - p_final = tf.reduce_sum(p_mult, 3) + p_nel_1D = fd.tfp_files.TruncatedSkewGaussianCC(loc=mean, scale=std_dev, + skewness=skew, + limit=ni_nel_2D, + owens_t_terms=owens_t_terms).prob(nel_2D) + + + #Restore p_nel unique_nel x nions-> unique_nel x n_electrons x n_photons x n_ions + p_nel=tf.gather_nd(params=p_nel_1D,indices=index_nel[:,o],batch_dims=0) + p_nel=tf.reshape(tf.reshape(p_nel,[-1]),[tf.shape(nq)[0],tf.shape(nq)[1],tf.shape(nq)[3]]) + p_nel=tf.repeat(p_nel[:,:,o,:],tf.shape(nq)[2],axis=2) + + #modified contractions remove need for costly repeats in ions dimension. + if self.is_ER: + p_mult = p_ni * p_nel + p_final = tf.reduce_sum(p_mult, 3)*p_nq + else: + p_mult = p_nq*p_nel + p_final = tf.tensordot(p_mult,p_ni_1D,axes=[[3],[0]]) r_final = p_final * rate_vs_energy @@ -169,6 +198,12 @@ def compute_single_energy_approx(args): return compute_single_energy(args, approx=True) nq = electrons_produced + photons_produced + #remove degenerate dimensions + #nevtxnelxnph->nq + unique_quanta,index_nq=tf.unique(tf.reshape(nq[:,:,:,0],[-1])) + #nevtxnel->nel' + unique_nel,index_nel=tf.unique(tf.reshape(electrons_produced[:,:,0,0],[-1])) + ions_min_initial = self.source._fetch('ions_produced_min', data_tensor=data_tensor)[:, 0, o] ions_min_initial = tf.repeat(ions_min_initial, tf.shape(ions_produced)[1], axis=1) @@ -179,6 +214,7 @@ def compute_single_energy_approx(args): # for the lowest energy ions_produced_add = ions_produced - ions_min_initial + # Energy above which we use the approximate computation if self.is_ER: cutoff_energy = 5. @@ -249,8 +285,16 @@ def _simulate(self, d): d['ions_produced'] = np.where(ni_temp < 0, ni_temp * 0, ni_temp) - - nex_temp = np.round(stats.norm.rvs(nq*alpha*ex_ratio, np.sqrt(nq*alpha*ex_ratio*nex_fano))).astype(int) + + mean_nex=nq*alpha*ex_ratio + width_nex=mean_nex*nex_fano + mean_nex=np.where(mean_nex < 0, + mean_nex * 0, + mean_nex) + width_nex=np.where(width_nex < 0, + width_nex * 0, + width_nex) + nex_temp = np.round(stats.norm.rvs(mean_nex, np.sqrt(width_nex))).astype(int) # Don't let number of excitons go negative nex = np.where(nex_temp < 0, nex_temp * 0, diff --git a/flamedisx/nest/lxe_sources.py b/flamedisx/nest/lxe_sources.py index 50aa8e037..ad78e60e5 100644 --- a/flamedisx/nest/lxe_sources.py +++ b/flamedisx/nest/lxe_sources.py @@ -22,8 +22,9 @@ class nestSource(fd.BlockModelSource): def __init__(self, *args, detector='default', **kwargs): + assert detector in ('default', 'lz','lz_WS2024') - + self.detector = detector assert os.path.exists(os.path.join( @@ -31,20 +32,18 @@ def __init__(self, *args, detector='default', **kwargs): config = configparser.ConfigParser(inline_comment_prefixes=';') config.read(os.path.join(os.path.dirname(__file__), 'config/', detector + '.ini')) - # common (known) parameters self.temperature = config.getfloat('NEST', 'temperature_config') self.pressure = config.getfloat('NEST', 'pressure_config') self.drift_field = config.getfloat('NEST', 'drift_field_config') self.gas_field = config.getfloat('NEST', 'gas_field_config') - # derived (known) parameters self.density = fd_nest.calculate_density( self.temperature, self.pressure) # NOTE: BE CAREFUL WITH THE BELOW, ONLY VALID NEAR VAPOUR PRESSURE!!! self.density_gas = fd_nest.calculate_density_gas( self.temperature, self.pressure) - # + self.drift_velocity = fd_nest.calculate_drift_velocity( self.drift_field, self.density, self.temperature, self.detector) self.Wq_keV, self.alpha = fd_nest.calculate_work(self.density) @@ -106,7 +105,6 @@ def recomb_prob(*args): nel_mean = args[0] nq_mean = args[1] ex_ratio = args[2] - elec_frac = nel_mean / nq_mean recomb_p = 1. - (ex_ratio + 1.) * elec_frac @@ -190,7 +188,7 @@ def s2_spe_smearing(self, n_pe): @export class nestERSource(nestSource): - def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, energy_bin_edges=None, **kwargs): + def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, energy_bin_edges=None,E_max_dim_size={'energy': 100}, **kwargs): if not hasattr(self, 'energies'): if energy_bin_edges is not None: self.energies = fd.np_to_tf(0.5 * (energy_bin_edges[1:] + energy_bin_edges[:-1])) @@ -199,6 +197,7 @@ def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, en self.energies = tf.cast(tf.linspace(energy_min, energy_max, num_energies), fd.float_type()) self.rates_vs_energy = tf.ones(num_energies, fd.float_type()) + self.max_dim_size=E_max_dim_size #edits dimension size of energy super().__init__(*args, **kwargs) model_blocks = ( @@ -216,25 +215,53 @@ def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, en # quanta_splitting.py - def mean_yield_electron(self, energy): + def mean_yield_electron(self, energy,*, + er_m1_a=30.66, + er_m1_b=6.1978, + er_m1_d=73.855, + er_m1_e=2.0318, + er_m1_f=0.41883, + + er_m10_a=0.0508273937, + er_m10_b=0.1166087199, + er_m10_c= 0.0508273937, + er_m10_d=1.39260460e+02, + er_m10_e=-0.65763592, + + er_Qy_a=77.2931084, + er_Qy_b=0.13946236, + er_Qy_c=0.52561312, + er_Qy_d=1.82217496, + er_Qy_e=2.82528809, + er_Qy_f=1.82217496, + er_Qy_g=144.65029656, + er_Qy_h=-2.80532006, + er_Qy_i=0.3344049589, + er_Qy_k=7.02921301, + er_Qy_l=98.27936794 , + er_Qy_m=7.0292130, + er_Qy_n=256.48156448, + er_Qy_o=1.29119251, + er_Qy_p=4.285781736 + ): Wq_eV = self.Wq_keV * 1e3 - Nq = energy * 1e3 / Wq_eV + Nq = energy * 1e3 / Wq_eV - m1 = tf.cast(30.66 + (6.1978 - 30.66) / pow(1. + pow(self.drift_field / 73.855, 2.0318), 0.41883), + m1 = tf.cast(er_m1_a + (er_m1_b - er_m1_a) / pow(1. + pow(self.drift_field / er_m1_d, er_m1_e), er_m1_f), fd.float_type()) m5 = tf.cast(Nq / energy / (1 + self.alpha * tf.math.erf(0.05 * energy)), fd.float_type()) - m1 - m10 = tf.cast((0.0508273937 + (0.1166087199 - 0.0508273937) / - (1 + pow(self.drift_field / 1.39260460e+02, -0.65763592))), + m10 = tf.cast((er_m10_a + (er_m10_b - er_m10_c) / + (1 + pow(self.drift_field / er_m10_d, er_m10_e))), fd.float_type()) - Qy = m1 + (77.2931084 - m1) / pow((1. + pow(energy / (fd.tf_log10(tf.cast(self.drift_field, fd.float_type())) * - 0.13946236 + 0.52561312), - 1.82217496 + (2.82528809 - 1.82217496) / - (1 + pow(self.drift_field / 144.65029656, -2.80532006)))), - 0.3344049589) + \ - m5 + (0. - m5) / pow((1. + pow(energy / (7.02921301 + (98.27936794 - 7.02921301) / - (1. + pow(self.drift_field / 256.48156448, 1.29119251))), 4.285781736)), m10) + Qy = m1 + (er_Qy_a - m1) / pow((1. + pow(energy / (fd.tf_log10(tf.cast(self.drift_field, fd.float_type())) * + er_Qy_b + er_Qy_c), + er_Qy_d + (er_Qy_e - er_Qy_f) / + (1 + pow(self.drift_field / er_Qy_g, er_Qy_h)))), + er_Qy_i) + \ + m5 + (0. - m5) / pow((1. + pow(energy / (er_Qy_k + (er_Qy_l - er_Qy_m) / + (1. + pow(self.drift_field / er_Qy_n, er_Qy_o))), er_Qy_p)), m10) coeff_TI = tf.cast(pow(1. / XENON_REF_DENSITY, 0.3), fd.float_type()) coeff_Ni = tf.cast(pow(1. / XENON_REF_DENSITY, 1.4), fd.float_type()) @@ -305,7 +332,6 @@ def skewness(self, nq_mean): if self.detector in ['lz','lz_WS2024']: skewness_masked = tf.zeros_like(nq_mean, dtype=fd.float_type()) - return skewness_masked def variance(self, *args): @@ -345,7 +371,7 @@ def variance(self, *args): @export class nestNRSource(nestSource): - def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, energy_bin_edges=None, **kwargs): + def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, energy_bin_edges=None,E_max_dim_size={'energy': 150}, **kwargs): if not hasattr(self, 'energies'): if energy_bin_edges is not None: self.energies = fd.np_to_tf(0.5 * (energy_bin_edges[1:] + energy_bin_edges[:-1])) @@ -354,6 +380,7 @@ def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, en self.energies = tf.cast(tf.linspace(energy_min, energy_max, num_energies), fd.float_type()) self.rates_vs_energy = tf.ones(num_energies, fd.float_type()) + self.max_dim_size=E_max_dim_size #edits dimension size of energy super().__init__(*args, **kwargs) model_blocks = ( @@ -371,23 +398,23 @@ def __init__(self, *args, energy_min=0.01, energy_max=10., num_energies=1000, en # quanta_splitting.py - def mean_yields(self, energy): - nr_nuis_a = 11. - nr_nuis_b = 1.1 - nr_nuis_c = 0.0480 - nr_nuis_d = -0.0533 - nr_nuis_e = 12.6 - nr_nuis_f = 0.3 - nr_nuis_g = 2. - nr_nuis_h = 0.3 - nr_nuis_i = 2 - nr_nuis_j = 0.5 - nr_nuis_k = 1. - nr_nuis_l = 1. - - TIB = nr_nuis_c * pow(self.drift_field, nr_nuis_d) * pow(self.density / XENON_REF_DENSITY, 0.3) - Qy = 1. / (TIB * pow(energy + nr_nuis_e, nr_nuis_j)) - Qy *= (1. - (1. / pow(1. + pow(energy / nr_nuis_f, nr_nuis_g), nr_nuis_k))) + def mean_yields(self, energy,* + ,nr_nuis_alpha = 11. + ,nr_nuis_beta = 1.1 + ,nr_nuis_gamma = 0.0480 + ,nr_nuis_delta = -0.0533 + ,nr_nuis_epsilon = 12.6 + ,nr_nuis_zeta = 0.3 + ,nr_nuis_eta = 2. + ,nr_nuis_theta = 0.3 + ,nr_nuis_l = 2 + ,nr_nuis_p = 0.5 + ,nr_new_nuis_a = 1. + ,nr_new_nuis_b = 1. + ): + TIB = nr_nuis_gamma * tf.math.pow(self.drift_field, nr_nuis_delta) * pow(self.density / XENON_REF_DENSITY, 0.3) + Qy = 1. / (TIB * tf.math.pow(energy + nr_nuis_epsilon, nr_nuis_p)) + Qy *= (1. - (1. / tf.math.pow(1. + tf.math.pow(tf.math.divide_no_nan(energy , nr_nuis_zeta), nr_nuis_eta), nr_new_nuis_a))) nel_temp = Qy * energy # Don't let number of electrons go negative @@ -395,9 +422,9 @@ def mean_yields(self, energy): 0 * nel_temp, nel_temp) - nq_temp = nr_nuis_a * pow(energy, nr_nuis_b) + nq_temp = nr_nuis_alpha * pow(energy, nr_nuis_beta) - nph_temp = (nq_temp - nel) * (1. - (1. / pow(1. + pow(energy / nr_nuis_h, nr_nuis_i), nr_nuis_l))) + nph_temp = (nq_temp - nel) * (1. - (1. / tf.math.pow(1. + tf.math.pow(tf.math.divide_no_nan(energy , nr_nuis_theta), nr_nuis_l), nr_new_nuis_b))) # Don't let number of photons go negative nph = tf.where(nph_temp < 0, 0 * nph_temp, @@ -409,7 +436,7 @@ def mean_yields(self, energy): nex = nq - ni - ex_ratio = tf.cast(nex / ni, fd.float_type()) + ex_ratio = tf.cast(tf.math.divide_no_nan(nex , ni), fd.float_type()) ex_ratio = tf.where(tf.logical_and(ex_ratio < self.alpha, energy > 100.), self.alpha * tf.ones_like(ex_ratio, dtype=fd.float_type()), @@ -528,7 +555,6 @@ def mean_yield_electron(self, energy): class nestFasterERSource(nestERSource): model_blocks = (fd_nest.FixedShapeEnergySpectrumFaster,) + nestERSource.model_blocks[1:] - @export class nestSpatialRateERSource(nestERSource): model_blocks = (fd_nest.SpatialRateEnergySpectrumER,) + nestERSource.model_blocks[1:] diff --git a/flamedisx/tfp_files/skew_gaussian.py b/flamedisx/tfp_files/skew_gaussian.py index c14fa833d..84e5df1c8 100644 --- a/flamedisx/tfp_files/skew_gaussian.py +++ b/flamedisx/tfp_files/skew_gaussian.py @@ -152,9 +152,10 @@ def owensT1(h, a, terms): exp_hs = tf.math.exp(hs) ci = -1 + exp_hs - val = tf.math.atan(a) * tf.ones_like(hs) - - for i in range(terms): + val = tf.math.atan(a) * tf.ones_like(hs)+ ci * a + ci= -ci + hs / tf.exp(tf.math.lgamma(tf.cast(2,'float32'))) * exp_hs + #for a->0. tf.math.pow(a,0.) can cause divergent gradients. Would be more efficeint not to calculate this. + for i in range(1,terms): val += ci * tf.math.pow(a,2*tf.cast(i,'float32')+1) / (2*tf.cast(i,'float32')+1) ci = -ci + tf.math.pow(hs,tf.cast(i+1,'float32')) / tf.exp(tf.math.lgamma(tf.cast(i+2,'float32'))) * exp_hs @@ -170,9 +171,9 @@ def _cdf(self, x): a = tf.cast(skewness,'float32') owens_t_eval = 0.5 * normal.Normal(loc=0.,scale=1.).cdf(h) + 0.5 * normal.Normal(loc=0.,scale=1.).cdf(a*h) - normal.Normal(loc=0.,scale=1.).cdf(h) * normal.Normal(loc=0.,scale=1.).cdf(a*h) - + #if tensorflow calculates a value, like 1/a., it will do the gradients as well. return 0.5 * (1. + tf.math.erf(1./(np.sqrt(2.)*scale) * (x - self.loc))) - \ - tf.cast(tf.where(a > tf.ones_like(a), 2. * (owens_t_eval - self.owensT1(a*h,1./a,self.owens_t_terms)), 2. * self.owensT1(h,a,self.owens_t_terms)),'float32') + tf.cast(tf.where(a > tf.ones_like(a), 2. * (owens_t_eval - self.owensT1(a*h,tf.math.divide_no_nan(1.,a),self.owens_t_terms)), 2. * self.owensT1(h,a,self.owens_t_terms)),'float32') def _parameter_control_dependencies(self, is_init): assertions = [] diff --git a/flamedisx/tfp_files/truncated_skew_gaussian_cc.py b/flamedisx/tfp_files/truncated_skew_gaussian_cc.py index f2d4ff009..d0275d290 100644 --- a/flamedisx/tfp_files/truncated_skew_gaussian_cc.py +++ b/flamedisx/tfp_files/truncated_skew_gaussian_cc.py @@ -136,7 +136,7 @@ def _event_shape_tensor(self): def _event_shape(self): return tf.TensorShape([]) - def _log_prob(self, x): + def prob(self, x): scale = tf.convert_to_tensor(self.scale) skewness = tf.convert_to_tensor(self.skewness) limit = tf.convert_to_tensor(self.limit) @@ -146,18 +146,18 @@ def _log_prob(self, x): cdf_lower = skew_gauss.cdf(x-0.5) minus_inf = dtype_util.as_numpy_dtype(x.dtype)(-np.inf) - - bounded_log_prob = tf.where((x > limit), - minus_inf, - tf.math.log(cdf_upper - cdf_lower)) - bounded_log_prob = tf.where(tf.math.is_nan(bounded_log_prob), - minus_inf, - bounded_log_prob) - dumping_log_prob = tf.where((x == limit), - tf.math.log(1 - cdf_lower), - bounded_log_prob) - - return dumping_log_prob + #_log_prob is used by exp(ln(x)), gradients are 1/x *e^ln(x). diverge at 0. + bounded_prob = tf.where((x > limit), + 0., + cdf_upper - cdf_lower+1e-11) + bounded_prob = tf.where(tf.math.is_nan(bounded_prob), + 0., + bounded_prob) + dumping_prob = tf.where((x == limit), + 1 - cdf_lower+1e-11, + bounded_prob) + + return dumping_prob def _parameter_control_dependencies(self, is_init): assertions = []