def trace_differential_rate(self): input_signature = (tf.TensorSpec(shape=self.data_tensor.shape[1:], dtype=fd.float_type()), tf.TensorSpec(shape=[len(self.defaults)], dtype=fd.float_type())) self._differential_rate_tf = tf.function( self._differential_rate, input_signature=input_signature)
def test_inference(xes: fd.ERSource): lf = fd.LogLikelihood( sources=dict(er=xes.__class__), elife=(100e3, 500e3, 5), data=xes.data) ## # Test non-autograph version ## x, x_grad = lf._log_likelihood(i_batch=tf.constant(0), dsetname=DEFAULT_DSETNAME, autograph=False, elife=tf.constant(200e3)) assert isinstance(x, tf.Tensor) assert x.dtype == fd.float_type() assert x.numpy() < 0 assert isinstance(x_grad, tf.Tensor) assert x_grad.dtype == fd.float_type() assert x_grad.numpy().shape == (1,) # Test a different parameter gives a different likelihood x2, x2_grad = lf._log_likelihood(i_batch=tf.constant(0), dsetname=DEFAULT_DSETNAME, autograph=False, elife=tf.constant(300e3)) assert (x - x2).numpy() != 0 assert (x_grad - x2_grad).numpy().sum() !=0 ## # Test batching # ## l1 = lf.log_likelihood(autograph=False) l2 = lf(autograph=False) lf.log_likelihood(elife=tf.constant(200e3), autograph=False)
def _populate_tensor_cache(self): super()._populate_tensor_cache() # Create an (n_time_bins, len(es)) histogram of spectra e_bin_centers = self.energy_hist.bin_centers(axis=1) e = np.array([ self.energy_hist.slicesum(t).histogram for t in self.data['t_j2000'] ]) # Look up in which time row/bin each event falls, and concatenate # the expected WIMP energy spectrum to the data tensor. # We modified _fetch so we can access these as 'wimp_energies' energy_tensor = tf.convert_to_tensor(e, dtype=fd.float_type()) assert energy_tensor.shape == [len(self.data), len(e_bin_centers)], \ f"{energy_tensor.shape} != {len(self.data)}, {len(e_bin_centers)}" energy_tensor = tf.reshape(energy_tensor, [self.n_batches, self.batch_size, -1]) self.data_tensor = tf.concat([self.data_tensor, energy_tensor], axis=2) # Store the centers of energy bins separately, these are the same # in each batch. es_centers = tf.convert_to_tensor(e_bin_centers, dtype=fd.float_type()) self.es_centers_batch = fd.repeat(es_centers[o, :], repeats=self.batch_size, axis=0)
def log_likelihood(self, autograph=True, second_order=False, omit_grads=tuple(), **kwargs): if second_order: # Compute the likelihood, jacobian and hessian # Use only non-tf.function version, in principle works with # but this leads to very long tracing times and we only need # hessian once f = self._log_likelihood_grad2 else: # Computes the likelihood and jacobian f = self._log_likelihood_tf if autograph else self._log_likelihood params = self.prepare_params(kwargs) n_grads = len(self.param_defaults) - len(omit_grads) ll = tf.constant(0., dtype=fd.float_type()) llgrad = tf.zeros(n_grads, dtype=fd.float_type()) llgrad2 = tf.zeros((n_grads, n_grads), dtype=fd.float_type()) for dsetname in self.dsetnames: for i_batch in tf.range(self.n_batches[dsetname], dtype=fd.int_type()): v = f(i_batch, dsetname, autograph, omit_grads=omit_grads, **params) ll += v[0] llgrad += v[1] if second_order: llgrad2 += v[2] if second_order: return ll, llgrad, llgrad2 return ll, llgrad
def p_electron_fluctuation(nq, q2=0.034, q3_nq=123. ): # From SR0, BBF model, right? # q3 = 1.7 keV ~= 123 quanta # For SR1: return tf.clip_by_value(q2 * (tf.constant(1.,dtype=fd.float_type()) - tf.exp(-nq / q3_nq)), tf.constant(1e-4,dtype=fd.float_type()), float('inf'))
def test_gimme(xes: fd.ERSource): x = xes.gimme('photon_gain_mean', data_tensor=None, ptensor=None) assert isinstance(x, tf.Tensor) assert x.dtype == fd.float_type() y = xes.gimme('photon_gain_mean', data_tensor=None, ptensor=None, numpy_out=True) assert isinstance(y, np.ndarray) if fd.float_type() == tf.float32: assert y.dtype == np.float32 else: assert y.dtype == np.float64 np.testing.assert_array_equal(x.numpy(), y) np.testing.assert_equal(y, xes.photon_gain_mean * np.ones(n_events)) data_tensor = xes.data_tensor[0] assert data_tensor is not None print(data_tensor.shape) z = xes.gimme('photon_gain_mean', data_tensor=data_tensor, ptensor=None) assert isinstance(z, tf.Tensor) assert z.dtype == fd.float_type() assert tf.reduce_all(tf.equal(x, z))
def trace_differential_rate(self): input_signature = (tf.TensorSpec(shape=self._batch_data_tensor_shape(), dtype=fd.float_type()), tf.TensorSpec(shape=[len(self.param_id)], dtype=fd.float_type())) self._differential_rate_tf = tf.function( self._differential_rate, input_signature=input_signature)
def test_inference(xes: fd.ERSource): lf = fd.LogLikelihood(sources=dict(er=xes.__class__), elife=(100e3, 500e3, 5), data=xes.data) # Test single-batch likelihood x, x_grad, _ = lf._log_likelihood( i_batch=tf.constant(0), dsetname=DEFAULT_DSETNAME, data_tensor=lf.data_tensors[DEFAULT_DSETNAME][0], batch_info=lf.batch_info, elife=tf.constant(200e3)) assert isinstance(x, tf.Tensor) assert x.dtype == fd.float_type() assert x.numpy() < 0 assert isinstance(x_grad, tf.Tensor) assert x_grad.dtype == fd.float_type() assert x_grad.numpy().shape == (1, ) # Test a different parameter gives a different likelihood x2, x2_grad, _ = lf._log_likelihood( i_batch=tf.constant(0), dsetname=DEFAULT_DSETNAME, data_tensor=lf.data_tensors[DEFAULT_DSETNAME][0], batch_info=lf.batch_info, elife=tf.constant(300e3)) assert (x - x2).numpy() != 0 assert (x_grad - x2_grad).numpy().sum() != 0 # Test batching l1 = lf.log_likelihood() l2 = lf() lf.log_likelihood(elife=tf.constant(200e3))
def s1_acceptance(s1, photon_detection_eff, photon_gain_mean, mean_eff=0.142 / (1 + 0.219)): # Both cS1 and S1 acceptance cs1 = mean_eff * s1 / (photon_detection_eff * photon_gain_mean) return tf.where((s1 < 2) | (s1 > 70) | (cs1 < 2), tf.zeros_like(s1, dtype=fd.float_type()), tf.ones_like(s1, dtype=fd.float_type()))
def gimme(self, fname, data_tensor=None, ptensor=None, bonus_arg=None, numpy_out=False): """Evaluate the model function fname with all required arguments :param fname: Name of the model function to compute :param bonus_arg: If fname takes a bonus argument, the data for it :param numpy_out: If True, return (tuple of) numpy arrays, otherwise (tuple of) tensors. :param data_tensor: Data tensor, columns as self.name_id If not given, use self.data (used in annotate) :param ptensor: Parameter tensor, columns as self.param_id If not give, use defaults dictionary (used in annotate) Before using gimme, you must use set_data to populate the internal caches. """ # TODO: make a clean way to keep track of i_batch or have it as input assert (bonus_arg is not None) == (fname in self.special_data_methods) if data_tensor is None: # We're in an annotate assert hasattr(self, 'data'), "You must set data first" else: # We're computing if not hasattr(self, 'name_id'): raise ValueError( "You must set_data first (and populate the tensor cache)") f = getattr(self, fname) if callable(f): args = [self._fetch(x, data_tensor) for x in self.f_dims[fname]] if bonus_arg is not None: args = [bonus_arg] + args kwargs = { pname: self._fetch_param(pname, ptensor) for pname in self.f_params[fname] } res = f(*args, **kwargs) else: if bonus_arg is None: n = len( self.data) if data_tensor is None else data_tensor.shape[0] x = tf.ones(n, dtype=fd.float_type()) else: x = tf.ones_like(bonus_arg, dtype=fd.float_type()) res = f * x if numpy_out: return fd.tf_to_np(res) return fd.np_to_tf(res)
def _populate_tensor_cache(self): super()._populate_tensor_cache() if self.spatial_rate_hist is not None: # Setup tensor of histogram for lookup positions = self.data[self.spatial_rate_hist_dims].values.T v = self.spatial_rate_hist.lookup(*positions) spatial_rate_tensor = tf.convert_to_tensor(v, dtype=fd.float_type()) self.spatial_rate_tensor = tf.reshape(spatial_rate_tensor, [self.n_batches, -1]) else: # If no hist defined, set uniform response self.spatial_rate_tensor = tf.ones( [self.n_batches, self.batch_size], dtype=fd.float_type())
def set_defaults(self, **params): for k, v in params.items(): if k in self.defaults: self.defaults[k] = tf.convert_to_tensor(v, dtype=fd.float_type()) else: raise ValueError(f"Key {k} not in defaults")
def p_el_sr0(e_kev): """Return probability of created quanta to become an electron for different deposited energies e_kev. This uses the charge yield model for XENON1T SR0, as published in https://arxiv.org/abs/1902.11297 (median posterior). """ e_kev = tf.convert_to_tensor(e_kev, dtype=fd.float_type()) # Parameters from Table II, for SR0 mean_nexni = 0.15 w_bbf = 13.8e-3 q0 = 1.13 q1 = 0.47 gamma_er = 0.124 / 4 omega_er = 31 f_dr = 120 delta_er = 0.24 with np.warnings.catch_warnings(): np.warnings.filterwarnings('ignore') fi = 1 / (1 + mean_nexni) nq = e_kev / w_bbf ni, nex = nq * fi, nq * (1 - fi) wiggle_er = gamma_er * tf.exp(-e_kev / omega_er) * f_dr ** (-delta_er) r_er = 1 - tf.math.log(1 + ni * wiggle_er) / (ni * wiggle_er) r_er /= (1 + tf.exp(-(e_kev - q0) / q1)) p_el = ni * (1 - r_er) / nq # placeholder value for e = 0 (better than NaN) p_el = tf.where(e_kev == 0, tf.ones_like(p_el), p_el) return p_el
def domain(self, x, data_tensor=None): """Return (n_events, |possible x values|) matrix containing all possible integer values of x for each event""" result1 = tf.cast(tf.range(self.dimsizes[x]), dtype=fd.float_type())[o, :] result2 = self._fetch(x + '_min', data_tensor=data_tensor)[:, o] return result1 + result2
def _log_likelihood_inner(self, i_batch, params, dsetname, autograph): # Does for loop over datasets and sources, not batches # Sum over sources is first in likelihood # Compute differential rates from all sources # drs = list[n_sources] of [n_events] tensors drs = tf.zeros((self.batch_size[dsetname],), dtype=fd.float_type()) for sname, s in self.sources.items(): if not self.d_for_s[sname] == dsetname: continue rate_mult = self._get_rate_mult(sname, params) dr = s.differential_rate(s.data_tensor[i_batch], autograph=autograph, **self._filter_source_kwargs(params, sname)) drs += dr * rate_mult # Sum over events and remove padding n = tf.where(tf.equal(i_batch, tf.constant(self.n_batches[dsetname] - 1, dtype=fd.int_type())), self.batch_size[dsetname] - self.n_padding[dsetname], self.batch_size[dsetname]) ll = tf.reduce_sum(tf.math.log(drs[:n])) # Add mu once (to the first batch) # and constraint really only once (to first batch of first dataset) ll += tf.where(tf.equal(i_batch, tf.constant(0, dtype=fd.int_type())), -self.mu(dsetname, **params) + (self.log_constraint(**params) if dsetname == self.dsetnames[0] else 0.), 0.) return ll
def detection_p(self, quanta_type, data_tensor, ptensor): """Return (n_events, |detected|, |produced|) tensor encoding P(n_detected | n_produced) """ n_det, n_prod = self.cross_domains(quanta_type + '_detected', quanta_type + '_produced', data_tensor) p = self.gimme(quanta_type + '_detection_eff', data_tensor=data_tensor, ptensor=ptensor)[:, o, o] if quanta_type == 'photon': # Note *= doesn't work, p will get reshaped p = p * self.gimme('penning_quenching_eff', bonus_arg=n_prod, data_tensor=data_tensor, ptensor=ptensor) result = tfp.distributions.Binomial( total_count=n_prod, probs=tf.cast(p, dtype=fd.float_type())).prob(n_det) return result * self.gimme(quanta_type + '_acceptance', bonus_arg=n_det, data_tensor=data_tensor, ptensor=ptensor)
def find_defaults(cls): """Discover which functions need which arguments / dimensions Discover possible parameters. Returns f_dims, f_params and defaults. """ f_dims = {x: [] for x in cls.data_methods} f_params = {x: [] for x in cls.data_methods} defaults = dict() for fname in cls.data_methods: f = getattr(cls, fname) if not callable(f): # Constant continue seen_special = False for pname, p in inspect.signature(f).parameters.items(): if pname == 'self': continue if p.default is inspect.Parameter.empty: if fname in cls.special_data_methods and not seen_special: seen_special = True else: # It's an observable dimension f_dims[fname].append(pname) else: # It's a parameter that can be fitted f_params[fname].append(pname) if (pname in defaults and p.default != defaults[pname]): raise ValueError(f"Inconsistent defaults for {pname}") defaults[pname] = tf.convert_to_tensor( p.default, dtype=fd.float_type()) return f_dims, f_params, defaults
def mu(self, dsetname, **kwargs): mu = tf.constant(0., dtype=fd.float_type()) for sname, s in self.sources.items(): if self.d_for_s[sname] != dsetname: continue mu += (self._get_rate_mult(sname, kwargs) * self.mu_itps[sname](**self._filter_source_kwargs(kwargs, sname))) return mu
def _populate_tensor_cache(self): super()._populate_tensor_cache() # Get energy bin centers e_bin_centers = self.energy_hist.bin_centers(axis=1) # Construct the energy spectra at event times e = np.array( [self.energy_hist.slicesum(t).histogram for t in self.data['t']]) energy_tensor = tf.convert_to_tensor(e, dtype=fd.float_type()) assert energy_tensor.shape == [len(self.data), len(e_bin_centers)], \ f"{energy_tensor.shape} != {len(self.data)}, {len(e_bin_centers)}" self.energy_tensor = tf.reshape(energy_tensor, [self.n_batches, self.batch_size, -1]) es_centers = tf.convert_to_tensor(e_bin_centers, dtype=fd.float_type()) self.all_es_centers = fd.repeat(es_centers[o, :], repeats=self.batch_size, axis=0)
def _populate_tensor_cache(self): # Create one big data tensor (n_batches, events_per_batch, n_cols) # TODO: make a list ctc = self.cols_to_cache self.data_tensor = tf.constant(self.data[ctc].values, dtype=fd.float_type()) self.data_tensor = tf.reshape( self.data_tensor, [self.n_batches, -1, len(ctc)])
def test_underscore_diff_rate(xes: fd.ERSource): x = xes._differential_rate(data_tensor=xes.data_tensor[0], ptensor=xes.ptensor_from_kwargs()) assert isinstance(x, tf.Tensor) assert x.dtype == fd.float_type() y = xes._differential_rate(data_tensor=xes.data_tensor[0], ptensor=xes.ptensor_from_kwargs(elife=100e3)) np.testing.assert_array_less(-fd.tf_to_np(tf.abs(x - y)), 0)
def rate_nq(self, nq_1d, data_tensor, ptensor): """Return differential rate at given number of produced quanta differs for ER and NR""" # TODO: this implementation echoes that for NR, but I feel there # must be a less clunky way... # (n_events, |ne|) tensors es, rate_e = self.gimme('energy_spectrum', data_tensor=data_tensor, ptensor=ptensor) q_produced = tf.cast(tf.floor(es / self.gimme( 'work', data_tensor=data_tensor, ptensor=ptensor)[:, o]), dtype=fd.float_type()) # (n_events, |nq|, |ne|) tensor giving p(nq | e) p_nq_e = tf.cast(tf.equal(nq_1d[:, :, o], q_produced[:, o, :]), dtype=fd.float_type()) q = tf.reduce_sum(p_nq_e * rate_e[:, o, :], axis=2) return q
def prepare_params(self, kwargs): for k in kwargs: if k not in self.param_defaults: raise ValueError(f"Unknown parameter {k}") # tf.function doesn't support {**x, **y} dict merging # return {**self.param_defaults, **kwargs} z = self.param_defaults.copy() for k, v in kwargs.items(): if isinstance(v, (float, int)) or fd.is_numpy_number(v): kwargs[k] = tf.constant(v, dtype=fd.float_type()) z.update(kwargs) return z
def mu_function(self, interpolation_method='star', n_trials=int(1e5), **param_specs): """Return interpolator for number of expected events Parameters must be specified as kwarg=(start, stop, n_anchors) """ if interpolation_method != 'star': raise NotImplementedError( f"mu interpolation method {interpolation_method} " f"not implemented") # Estimate mu under the current defaults base_mu = tf.constant(self.estimate_mu(n_trials=n_trials), dtype=fd.float_type()) # Estimate mus under the specified variations pspaces = dict() # parameter -> tf.linspace of anchors mus = dict() # parameter -> tensor of mus for pname, pspace_spec in tqdm(param_specs.items(), desc="Estimating mus"): pspaces[pname] = tf.linspace(*pspace_spec) mus[pname] = tf.convert_to_tensor([ self.estimate_mu(**{pname: x}, n_trials=n_trials) for x in np.linspace(*pspace_spec) ], dtype=fd.float_type()) def mu_itp(**kwargs): mu = base_mu for pname, v in kwargs.items(): mu *= tfp.math.interp_regular_1d_grid( x=v, x_ref_min=param_specs[pname][0], x_ref_max=param_specs[pname][1], y_ref=mus[pname]) / base_mu return mu return mu_itp
def lindhard_l(e, lindhard_k=tf.constant(0.138, dtype=fd.float_type())): """Return Lindhard quenching factor at energy e in keV""" eps = e * tf.constant(11.5 * 54.**(-7. / 3.), dtype=fd.float_type()) # Xenon: Z = 54 n0 = tf.constant(3., dtype=fd.float_type()) n1 = tf.constant(0.7, dtype=fd.float_type()) n2 = tf.constant(1.0, dtype=fd.float_type()) p0 = tf.constant(0.15, dtype=fd.float_type()) p1 = tf.constant(0.6, dtype=fd.float_type()) g = n0 * tf.pow(eps, p0) + n1 * tf.pow(eps, p1) + eps res = lindhard_k * g / (n2 + lindhard_k * g) return res
def rate_nphnel(self, data_tensor, ptensor): """Return differential rate tensor (n_events, |photons_produced|, |electrons_produced|) """ # Get differential rate and electron probability vs n_quanta # these four are (n_events, |nq|) tensors _nq_1d = self.domain('nq', data_tensor) rate_nq = self.rate_nq(_nq_1d, data_tensor=data_tensor, ptensor=ptensor) pel = self.gimme('p_electron', bonus_arg=_nq_1d, data_tensor=data_tensor, ptensor=ptensor) # Create tensors with the dimensions of our fin al result # i.e. (n_events, |photons_produced|, |electrons_produced|), # containing: # ... numbers of photons and electrons produced: nph, nel = self.cross_domains('photon_produced', 'electron_produced', data_tensor) # ... numbers of total quanta produced nq = nel + nph # ... indices in nq arrays _nq_ind = nq - self._fetch('nq_min', data_tensor=data_tensor)[:, o, o] # ... differential rate rate_nq = fd.lookup_axis1(rate_nq, _nq_ind) # ... probability of a quantum to become an electron pel = fd.lookup_axis1(pel, _nq_ind) # Finally, the main computation is simple: pel = tf.where(tf.math.is_nan(pel), tf.zeros_like(pel, dtype=fd.float_type()), pel) pel = tf.clip_by_value(pel, 1e-6, 1. - 1e-6) if self.do_pel_fluct: pel_fluct = self.gimme('p_electron_fluctuation', bonus_arg=_nq_1d, data_tensor=data_tensor, ptensor=ptensor) pel_fluct = fd.lookup_axis1(pel_fluct, _nq_ind) pel_fluct = tf.clip_by_value(pel_fluct, fd.MIN_FLUCTUATION_P, 1.) return rate_nq * fd.beta_binom_pmf( nph, n=nq, p_mean=1. - pel, p_sigma=pel_fluct) else: return rate_nq * tfp.distributions.Binomial(total_count=nq, probs=pel).prob(nel)
def p_electron(nq, W=13.8e-3, mean_nexni=0.135, q0=1.13, q1=0.47, gamma_er=0.031 , omega_er=31.): # gamma_er from paper 0.124/4 F = tf.constant(81.,dtype=fd.float_type()) e_kev = nq * W fi = 1. / (1. + mean_nexni) ni, nex = nq * fi, nq * (1. - fi) wiggle_er = gamma_er * tf.exp(-e_kev / omega_er) * F ** (-0.24) # delta_er and gamma_er are highly correlated # F **(-delta_er) set to constant r_er = 1. - tf.math.log(1. + ni * wiggle_er) / (ni * wiggle_er) r_er /= (1. + tf.exp(-(e_kev - q0) / q1)) p_el = ni * (1. - r_er) / nq return fd.safe_p(p_el)
def _log_likelihood_inner(self, i_batch, params, dsetname, data_tensor, batch_info): """Return log likelihood contribution of one batch in a dataset This loops over sources in the dataset and events in the batch, but not not over datasets or batches. """ # Retrieve batching info. Cannot use tuple-unpacking, tensorflow # doesn't like it when you iterate over tenstors dataset_index = self.dsetnames.index(dsetname) n_batches = batch_info[dataset_index, 0] batch_size = batch_info[dataset_index, 1] n_padding = batch_info[dataset_index, 2] # Compute differential rates from all sources # drs = list[n_sources] of [n_events] tensors drs = tf.zeros((batch_size, ), dtype=fd.float_type()) for source_i, sname in enumerate(self.sources_in_dset[dsetname]): s = self.sources[sname] rate_mult = self._get_rate_mult(sname, params) col_start, col_stop = self.column_indices[dsetname][source_i] dr = s.differential_rate( data_tensor[:, col_start:col_stop], # We are already tracing; if we call the traced function here # it breaks the Hessian (it will give NaNs) autograph=False, **self._filter_source_kwargs(params, sname)) drs += dr * rate_mult # Sum over events and remove padding n = tf.where(tf.equal(i_batch, n_batches - 1), batch_size - n_padding, batch_size) ll = tf.reduce_sum(tf.math.log(drs[:n])) # Add mu once (to the first batch) # and constraint really only once (to first batch of first dataset) ll += tf.where( tf.equal(i_batch, tf.constant(0, dtype=fd.int_type())), -self.mu(dsetname, **params) + (self.log_constraint( **params) if dsetname == self.dsetnames[0] else 0.), 0.) return ll
def _populate_tensor_cache(self): # Cache only float and int cols cols_to_cache = [ x for x in self.data.columns if fd.is_numpy_number(self.data[x]) ] self.name_id = tf.lookup.StaticVocabularyTable( tf.lookup.KeyValueTensorInitializer( tf.constant(cols_to_cache), tf.range(len(cols_to_cache), dtype=tf.dtypes.int64)), num_oov_buckets=1, lookup_key_dtype=tf.dtypes.string) # Create one big data tensor (n_batches, events_per_batch, n_cols) # TODO: make a list self.data_tensor = tf.constant(self.data[cols_to_cache].values, dtype=fd.float_type()) self.data_tensor = tf.reshape( self.data_tensor, [self.n_batches, -1, len(cols_to_cache)])
def test_hessian(xes: fd.ERSource): # Test the hessian at the guess position lf = fd.LogLikelihood( sources=dict(er=xes.__class__), elife=(100e3, 500e3, 5), free_rates='er', data=xes.data) guess = lf.guess() assert len(guess) == 2 inv_hess = lf.inverse_hessian(guess) inv_hess_np = inv_hess.numpy() assert inv_hess_np.shape == (2, 2) assert inv_hess.dtype == fd.float_type() # Check symmetry of hessian # The hessian is explicitly symmetrized before being passed to # the optimizer in bestfit a = inv_hess_np[0, 1] b = inv_hess_np[1, 0] assert abs(a - b)/(a+b) < 1e-3