From 16a51925028bc70aa6a9e590ad8fefd469b74a7d Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 22:25:33 -0500 Subject: [PATCH 01/32] Independent precision for loss module --- egenerator/loss/default.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/egenerator/loss/default.py b/egenerator/loss/default.py index 23dcaaea..7f9de612 100644 --- a/egenerator/loss/default.py +++ b/egenerator/loss/default.py @@ -192,9 +192,20 @@ def get_loss( "Sorting of loss terms is unnecessary when reducing to scalar" ) + # cast to specified float precision + precision = self.configuration.config["config"]["float_precision"] + data_batch_dict_cast = { + key: tf.cast(value, precision) + for key, value in data_batch_dict.items() + } + result_tensors_cast = { + key: tf.cast(value, precision) + for key, value in result_tensors.items() + } + loss_terms = self.loss_function( - data_batch_dict=data_batch_dict, - result_tensors=result_tensors, + data_batch_dict=data_batch_dict_cast, + result_tensors=result_tensors_cast, tensors=tensors, sort_loss_terms=sort_loss_terms, ) From a5983507017dfe135bfaf18e549dd1c750d6e629 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 22:32:33 -0500 Subject: [PATCH 02/32] Set loss precision to float64 --- configs/cascade_11param_noise_ftpv3m.yaml | 10 +++++----- configs/cascade_7param_noise_ftpv3m.yaml | 8 ++++---- configs/track_sphere_6param_ftpv3m.yaml | 10 +++++----- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/configs/cascade_11param_noise_ftpv3m.yaml b/configs/cascade_11param_noise_ftpv3m.yaml index 087356e3..cda06311 100644 --- a/configs/cascade_11param_noise_ftpv3m.yaml +++ b/configs/cascade_11param_noise_ftpv3m.yaml @@ -293,7 +293,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -307,7 +307,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -321,7 +321,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -335,7 +335,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -349,7 +349,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # define uniform priors 'uniform_parameters': { 'cascade_Absorption': [0.913, 1.087], diff --git a/configs/cascade_7param_noise_ftpv3m.yaml b/configs/cascade_7param_noise_ftpv3m.yaml index 6df2fb4d..9babdf61 100644 --- a/configs/cascade_7param_noise_ftpv3m.yaml +++ b/configs/cascade_7param_noise_ftpv3m.yaml @@ -286,7 +286,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -300,7 +300,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -314,7 +314,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -328,7 +328,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use diff --git a/configs/track_sphere_6param_ftpv3m.yaml b/configs/track_sphere_6param_ftpv3m.yaml index f15c8d6e..27b3ecd5 100644 --- a/configs/track_sphere_6param_ftpv3m.yaml +++ b/configs/track_sphere_6param_ftpv3m.yaml @@ -274,7 +274,7 @@ loss_module_settings: [ config: { # the float precision to use - 'float_precision': 'float32', + 'float_precision': 'float64', # Add normalization terms to llh if True 'add_normalization_term': True, # choose the loss function to use @@ -288,7 +288,7 @@ loss_module_settings: [ # config: { # # the float precision to use - # 'float_precision': 'float32', + # 'float_precision': 'float64', # # Add normalization terms to llh if True # 'add_normalization_term': True, # # choose the loss function to use @@ -302,7 +302,7 @@ loss_module_settings: [ # config: { # # the float precision to use - # 'float_precision': 'float32', + # 'float_precision': 'float64', # # Add normalization terms to llh if True # 'add_normalization_term': True, # # choose the loss function to use @@ -316,7 +316,7 @@ loss_module_settings: [ # config: { # # the float precision to use - # 'float_precision': 'float32', + # 'float_precision': 'float64', # # Add normalization terms to llh if True # 'add_normalization_term': True, # # choose the loss function to use @@ -330,7 +330,7 @@ loss_module_settings: [ # config: { # # the float precision to use - # 'float_precision': 'float32', + # 'float_precision': 'float64', # # Add normalization terms to llh if True # 'add_normalization_term': True, # # choose the loss function to use From 1eac373571456c76fff4257bc84fb43131c87e21 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:01:03 -0500 Subject: [PATCH 03/32] Add dtype argument for casting to basis functions --- egenerator/utils/basis_functions.py | 187 ++++++++++++++++++++++++---- 1 file changed, 164 insertions(+), 23 deletions(-) diff --git a/egenerator/utils/basis_functions.py b/egenerator/utils/basis_functions.py index 4a06058d..2c0b3443 100644 --- a/egenerator/utils/basis_functions.py +++ b/egenerator/utils/basis_functions.py @@ -5,7 +5,49 @@ from scipy.integrate import quad -def tf_gauss(x, mu, sigma): +def cast(dtype, *args): + """Cast all input arguments to the provided data type. + + Parameters + ---------- + dtype : str + The data type to cast the inputs to. + If None, no casting is performed. + *args : array_like + The input arguments to cast. + + Returns + ------- + array_like + The casted input arguments. + """ + if dtype is None: + return args + return tuple(np.array(arg, dtype=dtype) for arg in args) + + +def tf_cast(dtype, *args): + """Cast all input arguments to the provided data type. + + Parameters + ---------- + dtype : str + The data type to cast the inputs to. + If None, no casting is performed. + *args : tf.Tensor + The input arguments to cast. + + Returns + ------- + tf.Tensor + The casted input arguments. + """ + if dtype is None: + return args + return tuple(tf.cast(arg, dtype) for arg in args) + + +def tf_gauss(x, mu, sigma, dtype=None): """Gaussian PDF Parameters @@ -16,18 +58,22 @@ def tf_gauss(x, mu, sigma): Mu parameter of Gaussian. sigma : tf.Tensor Sigma parameter of Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The Gaussian PDF evaluated at x """ + x, mu, sigma = tf_cast(dtype, x, mu, sigma) return ( tf.exp(-0.5 * ((x - mu) / sigma) ** 2) / (2 * np.pi * sigma**2) ** 0.5 ) -def gauss(x, mu, sigma): +def gauss(x, mu, sigma, dtype=None): """Gaussian PDF Parameters @@ -38,18 +84,22 @@ def gauss(x, mu, sigma): Mu parameter of Gaussian. sigma : array_like Sigma parameter of Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The Gaussian PDF evaluated at x """ + x, mu, sigma = cast(dtype, x, mu, sigma) return ( np.exp(-0.5 * ((x - mu) / sigma) ** 2) / (2 * np.pi * sigma**2) ** 0.5 ) -def tf_log_gauss(x, mu, sigma): +def tf_log_gauss(x, mu, sigma, dtype=None): """Log Gaussian PDF Parameters @@ -60,17 +110,21 @@ def tf_log_gauss(x, mu, sigma): Mu parameter of Gaussian. sigma : tf.Tensor Sigma parameter of Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The Gaussian PDF evaluated at x """ + x, mu, sigma = tf_cast(dtype, x, mu, sigma) norm = np.log(np.sqrt(2 * np.pi)) return -0.5 * ((x - mu) / sigma) ** 2 - tf.math.log(sigma) - norm -def log_gauss(x, mu, sigma): +def log_gauss(x, mu, sigma, dtype=None): """Log Gaussian PDF Parameters @@ -81,17 +135,21 @@ def log_gauss(x, mu, sigma): Mu parameter of Gaussian. sigma : array_like Sigma parameter of Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The Gaussian PDF evaluated at x """ + x, mu, sigma = cast(dtype, x, mu, sigma) norm = np.log(np.sqrt(2 * np.pi)) return -0.5 * ((x - mu) / sigma) ** 2 - np.log(sigma) - norm -def tf_log_asymmetric_gauss(x, mu, sigma, r): +def tf_log_asymmetric_gauss(x, mu, sigma, r, dtype=None): """Asymmetric Log Gaussian PDF Parameters @@ -104,12 +162,16 @@ def tf_log_asymmetric_gauss(x, mu, sigma, r): Sigma parameter of Gaussian. r : tf.Tensor The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The asymmetric Gaussian PDF evaluated at x """ + x, mu, sigma, r = tf_cast(dtype, x, mu, sigma, r) norm = tf.math.log(2.0 / (tf.sqrt(2 * np.pi * sigma**2) * (r + 1))) exp = tf.where( x < mu, @@ -119,7 +181,7 @@ def tf_log_asymmetric_gauss(x, mu, sigma, r): return norm + exp -def log_asymmetric_gauss(x, mu, sigma, r): +def log_asymmetric_gauss(x, mu, sigma, r, dtype=None): """Asymmetric Log Gaussian PDF Parameters @@ -132,12 +194,16 @@ def log_asymmetric_gauss(x, mu, sigma, r): Sigma parameter of Gaussian. r : array_like The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The asymmetric Gaussian PDF evaluated at x """ + x, mu, sigma, r = cast(dtype, x, mu, sigma, r) norm = np.log(2.0 / (np.sqrt(2 * np.pi * sigma**2) * (r + 1))) exp = np.where( x < mu, @@ -147,7 +213,7 @@ def log_asymmetric_gauss(x, mu, sigma, r): return norm + exp -def tf_asymmetric_gauss(x, mu, sigma, r): +def tf_asymmetric_gauss(x, mu, sigma, r, dtype=None): """Asymmetric Gaussian PDF Parameters @@ -160,12 +226,16 @@ def tf_asymmetric_gauss(x, mu, sigma, r): Sigma parameter of Gaussian. r : tf.Tensor The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The asymmetric Gaussian PDF evaluated at x """ + x, mu, sigma, r = tf_cast(dtype, x, mu, sigma, r) norm = 2.0 / (tf.sqrt(2 * np.pi * sigma**2) * (r + 1)) exp = tf.where( x < mu, @@ -175,7 +245,7 @@ def tf_asymmetric_gauss(x, mu, sigma, r): return norm * exp -def asymmetric_gauss(x, mu, sigma, r): +def asymmetric_gauss(x, mu, sigma, r, dtype=None): """Asymmetric Gaussian PDF Parameters @@ -188,12 +258,16 @@ def asymmetric_gauss(x, mu, sigma, r): Sigma parameter of Gaussian. r : array_like The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The asymmetric Gaussian PDF evaluated at x """ + x, mu, sigma, r = cast(dtype, x, mu, sigma, r) norm = 2.0 / (np.sqrt(2 * np.pi * sigma**2) * (r + 1)) exp = np.where( x < mu, @@ -203,7 +277,7 @@ def asymmetric_gauss(x, mu, sigma, r): return norm * exp -def tf_asymmetric_gauss_cdf(x, mu, sigma, r): +def tf_asymmetric_gauss_cdf(x, mu, sigma, r, dtype=None): """Asymmetric Gaussian CDF Parameters @@ -216,12 +290,16 @@ def tf_asymmetric_gauss_cdf(x, mu, sigma, r): Sigma parameter of Gaussian. r : tf.Tensor The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The asymmetric Gaussian CDF evaluated at x """ + x, mu, sigma, r = tf_cast(dtype, x, mu, sigma, r) norm = 1.0 / (r + 1) exp = tf.where( x < mu, @@ -231,7 +309,7 @@ def tf_asymmetric_gauss_cdf(x, mu, sigma, r): return norm * exp -def asymmetric_gauss_cdf(x, mu, sigma, r): +def asymmetric_gauss_cdf(x, mu, sigma, r, dtype=None): """Asymmetric Gaussian CDF Parameters @@ -244,12 +322,16 @@ def asymmetric_gauss_cdf(x, mu, sigma, r): Sigma parameter of Gaussian. r : array_like The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The asymmetric Gaussian CDF evaluated at x """ + x, mu, sigma, r = cast(dtype, x, mu, sigma, r) norm = 1.0 / (r + 1) exp = np.where( x < mu, @@ -259,7 +341,7 @@ def asymmetric_gauss_cdf(x, mu, sigma, r): return norm * exp -def tf_asymmetric_gauss_ppf(q, mu, sigma, r): +def tf_asymmetric_gauss_ppf(q, mu, sigma, r, dtype=None): """Asymmetric Gaussian PPF Parameters @@ -272,12 +354,16 @@ def tf_asymmetric_gauss_ppf(q, mu, sigma, r): Sigma parameter of Gaussian. r : tf.Tensor The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The asymmetric Gaussian PPF evaluated at q """ + q, mu, sigma, r = tf_cast(dtype, q, mu, sigma, r) return tf.where( q < 1.0 / (r + 1), mu + np.sqrt(2) * sigma * tf.math.erfinv(q * (r + 1) - 1), @@ -285,7 +371,7 @@ def tf_asymmetric_gauss_ppf(q, mu, sigma, r): ) -def asymmetric_gauss_ppf(q, mu, sigma, r): +def asymmetric_gauss_ppf(q, mu, sigma, r, dtype=None): """Asymmetric Gaussian PPF Parameters @@ -298,12 +384,16 @@ def asymmetric_gauss_ppf(q, mu, sigma, r): Sigma parameter of Gaussian. r : array_like The asymmetry of the Gaussian. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The asymmetric Gaussian PPF evaluated at q """ + q, mu, sigma, r = cast(dtype, q, mu, sigma, r) return np.where( q < 1.0 / (r + 1), mu + np.sqrt(2) * sigma * special.erfinv(q * (r + 1) - 1), @@ -311,7 +401,9 @@ def asymmetric_gauss_ppf(q, mu, sigma, r): ) -def tf_log_negative_binomial(x, mu, alpha, add_normalization_term=False): +def tf_log_negative_binomial( + x, mu, alpha, add_normalization_term=False, dtype=None +): """Computes the logarithm of the negative binomial PDF The parameterization chosen here is defined by the mean mu and @@ -339,12 +431,17 @@ def tf_log_negative_binomial(x, mu, alpha, add_normalization_term=False): binomial distribution only has a proper normalization for integer x. For real-valued x the negative binomial is not properly normalized and hence adding the normalization term does not help. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor Logarithm of the negative binomal PDF evaluated at x. """ + x, mu, alpha = tf_cast(dtype, x, mu, alpha) + inv_alpha = 1.0 / alpha alpha_mu = alpha * mu @@ -360,7 +457,9 @@ def tf_log_negative_binomial(x, mu, alpha, add_normalization_term=False): return gamma_terms + term1 + term2 -def log_negative_binomial(x, mu, alpha, add_normalization_term=False): +def log_negative_binomial( + x, mu, alpha, add_normalization_term=False, dtype=None +): """Computes the logarithm of the negative binomial PDF The parameterization chosen here is defined by the mean mu and @@ -395,12 +494,17 @@ def negative_binomial(x, mu, alpha): binomial distribution only has a proper normalization for integer x. For real-valued x the negative binomial is not properly normalized and hence adding the normalization term does not help. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like Logarithm of the negative binomal PDF evaluated at x. """ + x, mu, alpha = cast(dtype, x, mu, alpha) + inv_alpha = 1.0 / alpha alpha_mu = alpha * mu @@ -508,7 +612,7 @@ def sample_from_negative_binomial( return rng.negative_binomial(r, p, size=size) -def negative_binomial_cdf(x, mu, alpha_or_var, param_is_alpha): +def negative_binomial_cdf(x, mu, alpha_or_var, param_is_alpha, dtype=None): """Computes the CDF of the negative binomial PDF The parameterization chosen here is defined by the mean mu and @@ -537,19 +641,23 @@ def negative_binomial_cdf(x, mu, alpha_or_var, param_is_alpha): param_is_alpha : bool If True, the parameter passed as `alpha_or_var` is alpha. If False, the parameter passed as `alpha_or_var` is the variance. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like CDF of the negative binomal PDF evaluated at x. """ + x, mu, alpha_or_var = cast(dtype, x, mu, alpha_or_var) p, r = convert_neg_binomial_params( mu=mu, alpha_or_var=alpha_or_var, param_is_alpha=param_is_alpha ) return stats.nbinom(r, p).cdf(x) -def negative_binomial_ppf(q, mu, alpha_or_var, param_is_alpha): +def negative_binomial_ppf(q, mu, alpha_or_var, param_is_alpha, dtype=None): """Computes the PPF of the negative binomial PDF The parameterization chosen here is defined by the mean mu and @@ -578,19 +686,23 @@ def negative_binomial_ppf(q, mu, alpha_or_var, param_is_alpha): param_is_alpha : bool If True, the parameter passed as `alpha_or_var` is alpha. If False, the parameter passed as `alpha_or_var` is the variance. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like PPF of the negative binomal PDF evaluated at x. """ + mu, alpha_or_var = cast(dtype, mu, alpha_or_var) p, r = convert_neg_binomial_params( mu=mu, alpha_or_var=alpha_or_var, param_is_alpha=param_is_alpha ) return stats.nbinom(r, p).ppf(q) -def tf_rayleigh(x, sigma): +def tf_rayleigh(x, sigma, dtype=None): """Computes Rayleigh PDF Parameters @@ -599,16 +711,20 @@ def tf_rayleigh(x, sigma): The input tensor. sigma : tf.Tensor The sigma parameter of the Rayleigh distribution. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The PDF of the Rayleigh distribution evaluated at x. """ + x, sigma = tf_cast(dtype, x, sigma) return x / (sigma**2) * tf.exp(-0.5 * (x / sigma) ** 2) -def rayleigh(x, sigma): +def rayleigh(x, sigma, dtype=None): """Computes Rayleigh PDF Parameters @@ -617,16 +733,20 @@ def rayleigh(x, sigma): The input tensor. sigma : array_like The sigma parameter of the Rayleigh distribution. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The PDF of the Rayleigh distribution evaluated at x. """ + x, sigma = cast(dtype, x, sigma) return x / (sigma**2) * np.exp(-0.5 * (x / sigma) ** 2) -def tf_rayleigh_cdf(x, sigma): +def tf_rayleigh_cdf(x, sigma, dtype=None): """Computes CDF of Rayleigh distribution. Parameters @@ -635,16 +755,20 @@ def tf_rayleigh_cdf(x, sigma): The input tensor. sigma : tf.Tensor The sigma parameter of the Rayleigh distribution. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- tf.Tensor The CDF of the Rayleigh distribution evaluated at x. """ + x, sigma = tf_cast(dtype, x, sigma) return 1 - tf.exp(-0.5 * (x / sigma) ** 2) -def rayleigh_cdf(x, sigma): +def rayleigh_cdf(x, sigma, dtype=None): """Computes CDF of Rayleigh distribution. Parameters @@ -653,16 +777,20 @@ def rayleigh_cdf(x, sigma): The input tensor. sigma : array_like The sigma parameter of the Rayleigh distribution. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The CDF of the Rayleigh distribution evaluated at x. """ + x, sigma = cast(dtype, x, sigma) return 1 - np.exp(-0.5 * (x / sigma) ** 2) -def von_mises_pdf(x, sigma, kent_min=np.deg2rad(7)): +def von_mises_pdf(x, sigma, kent_min=np.deg2rad(7), dtype=None): """Computes the von Mises-Fisher PDF on the sphere The PDF is defined and normalized in cartesian @@ -679,12 +807,16 @@ def von_mises_pdf(x, sigma, kent_min=np.deg2rad(7)): The value over which to use the von Mises-Fisher distribution. Underneath, a 2D Gaussian approximation is used for more numerical stability. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The PDF evaluated at the provided opening angles x. """ + x, sigma = cast(dtype, x, sigma) x = np.atleast_1d(x) sigma = np.atleast_1d(sigma) @@ -708,7 +840,7 @@ def von_mises_pdf(x, sigma, kent_min=np.deg2rad(7)): return result -def von_mises_in_dPsi_pdf(x, sigma, kent_min=np.deg2rad(7)): +def von_mises_in_dPsi_pdf(x, sigma, kent_min=np.deg2rad(7), dtype=None): """Computes the von Mises-Fisher PDF on the sphere The PDF is defined and normalized in the opening angle dPsi. @@ -723,12 +855,17 @@ def von_mises_in_dPsi_pdf(x, sigma, kent_min=np.deg2rad(7)): The value over which to use the von Mises-Fisher distribution. Underneath, a 2D Gaussian approximation is used for more numerical stability. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- array_like The PDF evaluated at the provided opening angles x. """ + x, sigma = cast(dtype, x, sigma) + # switching coordinates from (dx1, dx2) to spherical # coordinates (dPsi, phi) means that we have to include # the jakobi determinant sin dPsi @@ -741,7 +878,7 @@ def von_mises_in_dPsi_pdf(x, sigma, kent_min=np.deg2rad(7)): ) -def von_mises_in_dPsi_cdf(x, sigma, kent_min=np.deg2rad(7)): +def von_mises_in_dPsi_cdf(x, sigma, kent_min=np.deg2rad(7), dtype=None): """Computes the von Mises-Fisher CDF on the sphere The underlying PDF is defined and normalized in the opening angle dPsi. @@ -758,6 +895,9 @@ def von_mises_in_dPsi_cdf(x, sigma, kent_min=np.deg2rad(7)): The value over which to use the von Mises-Fisher distribution. Underneath, a 2D Gaussian approximation is used for more numerical stability. + dtype : str, optional + The data type of the output tensor, by default None. + If provided, the inputs are cast to this data type. Returns ------- @@ -765,6 +905,7 @@ def von_mises_in_dPsi_cdf(x, sigma, kent_min=np.deg2rad(7)): The CDF evaluated at the provided opening angles x. """ assert len(x) == len(sigma), ("Unequal lengths:", len(x), len(sigma)) + x, sigma = cast(dtype, x, sigma) x = np.atleast_1d(x) sigma = np.atleast_1d(sigma) result = [] From b42c22f4f1f797f00f59d9bea68775b040edc976 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:11:49 -0500 Subject: [PATCH 04/32] Only cast for float types --- egenerator/loss/default.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/egenerator/loss/default.py b/egenerator/loss/default.py index 7f9de612..97a64c53 100644 --- a/egenerator/loss/default.py +++ b/egenerator/loss/default.py @@ -197,10 +197,12 @@ def get_loss( data_batch_dict_cast = { key: tf.cast(value, precision) for key, value in data_batch_dict.items() + if value.dtype in (tf.float16, tf.float32, tf.float64) } result_tensors_cast = { key: tf.cast(value, precision) for key, value in result_tensors.items() + if value.dtype in (tf.float16, tf.float32, tf.float64) } loss_terms = self.loss_function( From ce8e9de6fe16359efb54b8007c918812681c4e49 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:20:51 -0500 Subject: [PATCH 05/32] Only cast for float types --- egenerator/loss/default.py | 6 ++++-- egenerator/loss/snowstorm.py | 5 +++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/egenerator/loss/default.py b/egenerator/loss/default.py index 97a64c53..9bee9dab 100644 --- a/egenerator/loss/default.py +++ b/egenerator/loss/default.py @@ -197,12 +197,14 @@ def get_loss( data_batch_dict_cast = { key: tf.cast(value, precision) for key, value in data_batch_dict.items() - if value.dtype in (tf.float16, tf.float32, tf.float64) + if tf.is_tensor(value) + and value.dtype in (tf.float16, tf.float32, tf.float64) } result_tensors_cast = { key: tf.cast(value, precision) for key, value in result_tensors.items() - if value.dtype in (tf.float16, tf.float32, tf.float64) + if tf.is_tensor(value) + and value.dtype in (tf.float16, tf.float32, tf.float64) } loss_terms = self.loss_function( diff --git a/egenerator/loss/snowstorm.py b/egenerator/loss/snowstorm.py index 42e6679f..d577dc30 100644 --- a/egenerator/loss/snowstorm.py +++ b/egenerator/loss/snowstorm.py @@ -195,6 +195,10 @@ def uniform_log_prior_loss(self, values, low, high, eps=0): normalization = np.exp(exp_factor) def loss_excess(scaled_excess): + scaled_excess = tf.cast( + scaled_excess, + self.configuration.config["config"]["float_precision"], + ) return tf.exp((scaled_excess + 1) * exp_factor) - normalization loss = tf.where( @@ -314,6 +318,7 @@ def get_loss( fourier_values, mu=tf.zeros_like(fourier_sigmas), sigma=fourier_sigmas, + dtype=self.configuration.config["config"]["float_precision"], ) # we will use the negative log likelihood as loss From 57c39f77a28af2e9ae41564895178f22e4d865d2 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:25:00 -0500 Subject: [PATCH 06/32] Only cast for float types --- egenerator/loss/default.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/egenerator/loss/default.py b/egenerator/loss/default.py index 9bee9dab..260af04c 100644 --- a/egenerator/loss/default.py +++ b/egenerator/loss/default.py @@ -194,18 +194,27 @@ def get_loss( # cast to specified float precision precision = self.configuration.config["config"]["float_precision"] - data_batch_dict_cast = { - key: tf.cast(value, precision) - for key, value in data_batch_dict.items() - if tf.is_tensor(value) - and value.dtype in (tf.float16, tf.float32, tf.float64) - } - result_tensors_cast = { - key: tf.cast(value, precision) - for key, value in result_tensors.items() - if tf.is_tensor(value) - and value.dtype in (tf.float16, tf.float32, tf.float64) - } + data_batch_dict_cast = {} + for key, value in data_batch_dict.items(): + if tf.is_tensor(value) and value.dtype in ( + tf.float16, + tf.float32, + tf.float64, + ): + data_batch_dict_cast[key] = tf.cast(value, precision) + else: + data_batch_dict_cast[key] = value + + result_tensors_cast = {} + for key, value in result_tensors.items(): + if tf.is_tensor(value) and value.dtype in ( + tf.float16, + tf.float32, + tf.float64, + ): + result_tensors_cast[key] = tf.cast(value, precision) + else: + result_tensors_cast[key] = value loss_terms = self.loss_function( data_batch_dict=data_batch_dict_cast, From 288bc3e9713bd91e2966186fe09350aa09d4a892 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:36:45 -0500 Subject: [PATCH 07/32] Only cast for float types --- egenerator/loss/default.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/egenerator/loss/default.py b/egenerator/loss/default.py index 260af04c..1bd9a827 100644 --- a/egenerator/loss/default.py +++ b/egenerator/loss/default.py @@ -236,10 +236,13 @@ def get_loss( loss_terms[2] = tf.zeros_like(dom_tensor) if normalize_by_total_charge: - total_charge = tf.clip_by_value( - tf.reduce_sum(data_batch_dict["x_pulses"][:, 0]), - 1, - float("inf"), + total_charge = tf.cast( + tf.clip_by_value( + tf.reduce_sum(data_batch_dict["x_pulses"][:, 0]), + 1, + float("inf"), + ), + dtype=precision, ) loss_terms = [loss / total_charge for loss in loss_terms] From 9613c7e9afe9cc923ecd6c2458709c477e780e35 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:51:57 -0500 Subject: [PATCH 08/32] Add dtype argument to numpy pdf and cdf functions --- egenerator/model/source/base.py | 34 ++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/egenerator/model/source/base.py b/egenerator/model/source/base.py index aff796f3..2e467ebd 100644 --- a/egenerator/model/source/base.py +++ b/egenerator/model/source/base.py @@ -405,6 +405,7 @@ def cdf( tw_exclusions_ids=None, strings=slice(None), doms=slice(None), + dtype="float64", **kwargs ): """Compute CDF values at x for given result_tensors @@ -457,6 +458,9 @@ def cdf( The doms to slice the PDF for. If None, all doms are used. Shape: [n_doms] + dtype : str, optional + The data type of the output array. + Default: 'float64' **kwargs Keyword arguments. @@ -474,7 +478,12 @@ def cdf( If asymmetric Gaussian latent variables are not present in `result_tensors` dictionary. """ - eps = 1e-7 + if dtype == "float64": + eps = 1e-15 + elif dtype == "float32": + eps = 1e-7 + else: + raise ValueError(f"Invalid dtype: {dtype}") x_orig = np.atleast_1d(x) assert len(x_orig.shape) == 1, x_orig.shape @@ -521,7 +530,11 @@ def cdf( # shape: [n_events, n_strings, n_doms, n_components, n_points] mixture_cdf = basis_functions.asymmetric_gauss_cdf( - x=x, mu=mu, sigma=sigma, r=r + x=x, + mu=mu, + sigma=sigma, + r=r, + dtype=dtype, ) # uniformly scale up pdf values due to excluded regions @@ -573,6 +586,7 @@ def cdf( sigma[ids[0], ids[1], ids[2]], [1, 1, -1] ), r=np.reshape(r[ids[0], ids[1], ids[2]], [1, 1, -1]), + dtype=dtype, ) * np.reshape(scale[ids[0], ids[1], ids[2]], [1, 1, -1]), axis=2, @@ -585,7 +599,7 @@ def cdf( cdf_values[ids[0], ids[1], ids[2]] -= cdf_excluded - eps = 1e-3 + eps = 1e-6 if (cdf_values < 0 - eps).any(): self._logger.warning( "CDF values below zero: {}".format( @@ -609,6 +623,7 @@ def pdf( tw_exclusions_ids=None, strings=slice(None), doms=slice(None), + dtype="float64", **kwargs ): """Compute PDF values at x for given result_tensors @@ -678,7 +693,12 @@ def pdf( If asymmetric Gaussian latent variables are not present in `result_tensors` dictionary. """ - eps = 1e-7 + if dtype == "float64": + eps = 1e-15 + elif dtype == "float32": + eps = 1e-7 + else: + raise ValueError(f"Invalid dtype: {dtype}") x_orig = np.atleast_1d(x) assert len(x_orig.shape) == 1, x_orig.shape @@ -722,7 +742,11 @@ def pdf( # shape: [n_events, n_strings, n_doms, n_components, n_points] mixture_pdf = basis_functions.asymmetric_gauss( - x=x, mu=mu, sigma=sigma, r=r + x=x, + mu=mu, + sigma=sigma, + r=r, + dtype=dtype, ) # uniformly scale up pdf values due to excluded regions From 89072c582f21a1547252ce857aa02b466e34319f Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:52:26 -0500 Subject: [PATCH 09/32] Force evaluation of PDF/CDF in float64 --- .../model/source/track/inf_tracks_sphere.py | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index af2d6bbc..3e9c05ec 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -694,18 +694,23 @@ def get_tensors( mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, + dtype="float64", ) tw_cdf_stop = basis_functions.tf_asymmetric_gauss_cdf( x=t_exclusions[:, 1], mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, + dtype="float64", ) # shape: [n_tw, n_models] tw_cdf_exclusion = tf_helpers.safe_cdf_clip( tw_cdf_stop - tw_cdf_start ) + tw_cdf_exclusion = tf.cast( + tw_cdf_exclusion, config["float_precision"] + ) # shape: [n_tw, 1] tw_cdf_exclusion_reduced = tf.reduce_sum( @@ -843,14 +848,18 @@ def get_tensors( # shape: [n_batch, 86, 60, 1] dom_charges_llh = tf.where( dom_charges_true > charge_threshold, - tf.math.log( - basis_functions.tf_asymmetric_gauss( - x=dom_charges_true, - mu=dom_charges, - sigma=dom_charges_sigma, - r=dom_charges_r, - ) - + self.epsilon + tf.cast( + tf.math.log( + basis_functions.tf_asymmetric_gauss( + x=dom_charges_true, + mu=dom_charges, + sigma=dom_charges_sigma, + r=dom_charges_r, + dtype="float64", + ) + + self.epsilon + ), + config["float_precision"], ), dom_charges_true * tf.math.log(dom_charges + self.epsilon) - dom_charges, @@ -897,6 +906,10 @@ def get_tensors( x=dom_charges_true, mu=dom_charges, alpha=dom_charges_alpha, + dtype="float64", + ) + dom_charges_llh = tf.cast( + dom_charges_llh, config["float_precision"] ) # compute standard deviation @@ -955,6 +968,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, + dtype="float64", ) * pulse_latent_scale ) @@ -964,6 +978,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, + dtype="float64", ) * pulse_latent_scale ) @@ -972,6 +987,10 @@ def get_tensors( pulse_pdf_values = tf.reduce_sum(pulse_pdf_values, axis=-1) pulse_cdf_values = tf.reduce_sum(pulse_cdf_values, axis=-1) + # cast back to specified float precision + pulse_pdf_values = tf.cast(pulse_pdf_values, config["float_precision"]) + pulse_cdf_values = tf.cast(pulse_cdf_values, config["float_precision"]) + # scale up pulse pdf by time exclusions if needed if time_exclusions_exist: From 579dc371fa8a765f93c50eacaa095131d7a5c410 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 5 Oct 2024 23:58:34 -0500 Subject: [PATCH 10/32] Force evaluation of PDF/CDF in float64 --- egenerator/model/source/track/inf_tracks_sphere.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index 3e9c05ec..3237996b 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -961,6 +961,8 @@ def get_tensors( # Apply Asymmetric Gaussian Mixture Model # ------------------------------------------- + pulse_latent_scale = tf.cast(pulse_latent_scale, "float64") + # [n_pulses, 1] * [n_pulses, n_models] = [n_pulses, n_models] pulse_pdf_values = ( basis_functions.tf_asymmetric_gauss( From 30eeb8873372979eb09eceabb74774965fe59ea1 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sun, 6 Oct 2024 00:30:30 -0500 Subject: [PATCH 11/32] safe_log: clip instead of +eps in log calls --- egenerator/loss/default.py | 25 +++++----- egenerator/model/source/cascade/default.py | 5 +- .../model/source/track/inf_tracks_sphere.py | 5 +- egenerator/utils/tf_helpers.py | 48 +++++++++++++++++++ 4 files changed, 65 insertions(+), 18 deletions(-) diff --git a/egenerator/loss/default.py b/egenerator/loss/default.py index 1bd9a827..875a04de 100644 --- a/egenerator/loss/default.py +++ b/egenerator/loss/default.py @@ -4,6 +4,7 @@ from egenerator import misc from egenerator.utils import basis_functions from egenerator.manager.component import BaseComponent, Configuration +from egenerator.utils import tf_helpers class DefaultLossModule(BaseComponent): @@ -362,15 +363,16 @@ def unbinned_extended_pulse_llh( dom_charges_pred = dom_charges_pred * mask_valid # prevent log(zeros) issues - pulse_log_pdf_values = tf.math.log(pulse_pdf_values + self.epsilon) + pulse_log_pdf_values = tf_helpers.safe_log(pulse_pdf_values) # compute unbinned negative likelihood over pulse times with given # time pdf: -sum( charge_i * log(pdf_d(t_i)) ) time_log_likelihood = -pulse_charges * pulse_log_pdf_values # get poisson likelihood over total charge at a DOM for extendended LLH - llh_poisson = dom_charges_pred - dom_charges_true * tf.math.log( - dom_charges_pred + self.epsilon + llh_poisson = ( + dom_charges_pred + - dom_charges_true * tf_helpers.safe_log(dom_charges_pred) ) if sort_loss_terms: @@ -467,7 +469,7 @@ def unbinned_pulse_time_llh( ), "Model must deal with time exclusions!" # prevent log(zeros) issues - pulse_log_pdf_values = tf.math.log(pulse_pdf_values + self.epsilon) + pulse_log_pdf_values = tf_helpers.safe_log(pulse_pdf_values) # compute unbinned negative likelihood over pulse times with given # time pdf: -sum( charge_i * log(pdf_d(t_i)) ) @@ -613,12 +615,11 @@ def unbinned_pulse_time_mpe_llh( # charge_i * pdf_i(t_0)^c_0 * (1 - cdf_i(t_0))^(charge_i - c_0) # with t_0 and c_0 the first pulse time and charge at DOM i # Shape: [n_pulses_first] - eps = self.epsilon mpe_log_llh = ( - tf.math.log(dom_charges_true_pulses + eps) - + pulse_charge_first * tf.math.log(pulse_pdf_value_first + eps) - + (dom_charges_true_pulses - pulse_charge_first + eps) - * tf.math.log(1 - pulse_cdf_value_first + eps) + tf_helpers.safe_log(dom_charges_true_pulses) + + pulse_charge_first * tf_helpers.safe_log(pulse_pdf_value_first) + + (dom_charges_true_pulses - pulse_charge_first) + * tf_helpers.safe_log(1 - pulse_cdf_value_first) ) time_loss = -mpe_log_llh @@ -746,7 +747,7 @@ def unbinned_pulse_and_dom_charge_pdf( llh_charge = llh_charge * mask_valid # prevent log(zeros) issues - pulse_log_pdf_values = tf.math.log(pulse_pdf_values + self.epsilon) + pulse_log_pdf_values = tf_helpers.safe_log(pulse_pdf_values) # compute unbinned negative likelihood over pulse times with given # time pdf: -sum( charge_i * log(pdf_d(t_i)) ) @@ -861,7 +862,7 @@ def unbinned_charge_quantile_pdf( ) # prevent log(zeros) issues - pulse_log_pdf_values = tf.math.log(pulse_pdf_values + self.epsilon) + pulse_log_pdf_values = tf_helpers.safe_log(pulse_pdf_values) # compute unbinned negative likelihood over pulse times with given # time pdf: -sum( log(pdf_d(t_i)) ) @@ -1421,7 +1422,7 @@ def normalized_dom_charge_pdf( # shape: [n_batch, 86, 60] dom_pdf = hits_pred / (event_total + self.epsilon) - llh_dom = hits_true * tf.math.log(dom_pdf + self.epsilon) + llh_dom = hits_true * tf_helpers.safe_log(dom_pdf) if sort_loss_terms: loss_terms = [ diff --git a/egenerator/model/source/cascade/default.py b/egenerator/model/source/cascade/default.py index 7c7fbc98..3cb7544c 100644 --- a/egenerator/model/source/cascade/default.py +++ b/egenerator/model/source/cascade/default.py @@ -772,16 +772,15 @@ def get_tensors( # shape: [n_batch, 86, 60, 1] dom_charges_llh = tf.where( dom_charges_true > charge_threshold, - tf.math.log( + tf_helpers.safe_log( basis_functions.tf_asymmetric_gauss( x=dom_charges_true, mu=dom_charges, sigma=dom_charges_sigma, r=dom_charges_r, ) - + self.epsilon ), - dom_charges_true * tf.math.log(dom_charges + self.epsilon) + dom_charges_true * tf_helpers.safe_log(dom_charges) - dom_charges, ) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index 3237996b..23d4e88e 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -849,7 +849,7 @@ def get_tensors( dom_charges_llh = tf.where( dom_charges_true > charge_threshold, tf.cast( - tf.math.log( + tf_helpers.safe_log( basis_functions.tf_asymmetric_gauss( x=dom_charges_true, mu=dom_charges, @@ -857,11 +857,10 @@ def get_tensors( r=dom_charges_r, dtype="float64", ) - + self.epsilon ), config["float_precision"], ), - dom_charges_true * tf.math.log(dom_charges + self.epsilon) + dom_charges_true * tf_helpers.safe_log(dom_charges) - dom_charges, ) diff --git a/egenerator/utils/tf_helpers.py b/egenerator/utils/tf_helpers.py index cfc1882c..c30afe86 100644 --- a/egenerator/utils/tf_helpers.py +++ b/egenerator/utils/tf_helpers.py @@ -1,6 +1,54 @@ import tensorflow as tf +def clip_logits(logits, eps=None): + """Clip logits to avoid numerical instabilities + + Parameters + ---------- + logits : tf.Tensor + The logits to clip. + eps : float, optional + The epsilon value to clip to, by default None. + If None the value will be clipped to 1e-37 for float32 + and 1e-307 for float64. + + Returns + ------- + tf.Tensor + The clipped logits. + """ + if eps is None: + # Up to these values, the log returns finite values + if logits.dtype == tf.float32: + eps = 1e-37 + elif logits.dtype == tf.float64: + eps = 1e-307 + else: + raise ValueError(f"Unknown dtype for logits: {logits.dtype}") + return tf.clip_by_value(logits, eps, float("inf")) + + +def safe_log(logits, eps=None): + """Safe log operation + + Parameters + ---------- + logits : tf.Tensor + The logits to log. + eps : float, optional + The epsilon value to clip to, by default None. + If None the value will be clipped to 1e-37 for float32 + and 1e-307 for float64. + + Returns + ------- + tf.Tensor + The safe log of the logits. + """ + return tf.math.log(clip_logits(logits, eps=eps)) + + def safe_cdf_clip(cdf_values, tol=1e-5): """Perform clipping of CDF values From eaf8570bfedb3e81f63292316e8e0c87cf56d6ad Mon Sep 17 00:00:00 2001 From: mhuen Date: Sun, 6 Oct 2024 21:47:03 -0500 Subject: [PATCH 12/32] Cascades: safe log and float64 pdf/cdf evaluation --- egenerator/model/source/cascade/default.py | 49 +++++++++++++++++----- 1 file changed, 38 insertions(+), 11 deletions(-) diff --git a/egenerator/model/source/cascade/default.py b/egenerator/model/source/cascade/default.py index 3cb7544c..2a684e23 100644 --- a/egenerator/model/source/cascade/default.py +++ b/egenerator/model/source/cascade/default.py @@ -621,18 +621,23 @@ def get_tensors( mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, + dtype="float64", ) tw_cdf_stop = basis_functions.tf_asymmetric_gauss_cdf( x=t_exclusions[:, 1], mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, + dtype="float64", ) # shape: [n_tw, n_models] tw_cdf_exclusion = tf_helpers.safe_cdf_clip( tw_cdf_stop - tw_cdf_start ) + tw_cdf_exclusion = tf.cast( + tw_cdf_exclusion, dtype=config["float_precision"] + ) # shape: [n_tw, 1] tw_cdf_exclusion_reduced = tf.reduce_sum( @@ -772,13 +777,17 @@ def get_tensors( # shape: [n_batch, 86, 60, 1] dom_charges_llh = tf.where( dom_charges_true > charge_threshold, - tf_helpers.safe_log( - basis_functions.tf_asymmetric_gauss( - x=dom_charges_true, - mu=dom_charges, - sigma=dom_charges_sigma, - r=dom_charges_r, - ) + tf.cast( + tf_helpers.safe_log( + basis_functions.tf_asymmetric_gauss( + x=dom_charges_true, + mu=dom_charges, + sigma=dom_charges_sigma, + r=dom_charges_r, + dtype="float64", + ) + ), + dtype=config["float_precision"], ), dom_charges_true * tf_helpers.safe_log(dom_charges) - dom_charges, @@ -821,10 +830,14 @@ def get_tensors( dom_charges_alpha = tf.nn.elu(alpha_trafo - 5) + 1.000001 # compute log pdf - dom_charges_llh = basis_functions.tf_log_negative_binomial( - x=dom_charges_true, - mu=dom_charges, - alpha=dom_charges_alpha, + dom_charges_llh = tf.cast( + basis_functions.tf_log_negative_binomial( + x=dom_charges_true, + mu=dom_charges, + alpha=dom_charges_alpha, + dtype="float64", + ), + dtype=config["float_precision"], ) # compute standard deviation @@ -876,6 +889,10 @@ def get_tensors( # Apply Asymmetric Gaussian Mixture Model # ------------------------------------------- + pulse_latent_scale = tf.cast( + pulse_latent_scale, dtype=config["float_precision"] + ) + # [n_pulses, 1] * [n_pulses, n_models] = [n_pulses, n_models] pulse_pdf_values = ( basis_functions.tf_asymmetric_gauss( @@ -883,6 +900,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, + dtype="float64", ) * pulse_latent_scale ) @@ -892,6 +910,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, + dtype="float64", ) * pulse_latent_scale ) @@ -900,6 +919,14 @@ def get_tensors( pulse_pdf_values = tf.reduce_sum(pulse_pdf_values, axis=-1) pulse_cdf_values = tf.reduce_sum(pulse_cdf_values, axis=-1) + # cast back to specified float precision + pulse_pdf_values = tf.cast( + pulse_pdf_values, dtype=config["float_precision"] + ) + pulse_cdf_values = tf.cast( + pulse_cdf_values, dtype=config["float_precision"] + ) + # scale up pulse pdf by time exclusions if needed if time_exclusions_exist: From dbf9e5a3793cbe69b15da2860834d6dbacc8b6d4 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sun, 6 Oct 2024 21:52:17 -0500 Subject: [PATCH 13/32] Revert boundaries --- egenerator/model/source/track/inf_tracks_sphere.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index 23d4e88e..f0ed9e0e 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -639,9 +639,9 @@ def get_tensors( latent_mu = dt_geometry + t_seed + factor_mu * latent_mu # force positive and min values - latent_scale = tf.math.exp(latent_scale) - latent_r = tf.math.exp(latent_r) - latent_sigma = tf.math.exp(latent_sigma) + 0.0001 + latent_scale = tf.nn.elu(latent_scale) + 1.00001 + latent_r = tf.nn.elu(latent_r) + 1.0001 + latent_sigma = tf.nn.elu(latent_sigma) + 1.0001 # normalize scale to sum to 1 latent_scale /= tf.reduce_sum(latent_scale, axis=-1, keepdims=True) From 798c2b9c545b47c7f71c5649d4802bd4717a18fb Mon Sep 17 00:00:00 2001 From: mhuen Date: Mon, 7 Oct 2024 08:51:46 -0500 Subject: [PATCH 14/32] Lower bound for sigma and r --- egenerator/model/source/cascade/default.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/egenerator/model/source/cascade/default.py b/egenerator/model/source/cascade/default.py index 2a684e23..f45a7543 100644 --- a/egenerator/model/source/cascade/default.py +++ b/egenerator/model/source/cascade/default.py @@ -567,8 +567,8 @@ def get_tensors( # force positive and min values latent_scale = tf.nn.elu(latent_scale) + 1.00001 - latent_r = tf.nn.elu(latent_r) + 1.001 - latent_sigma = tf.nn.elu(latent_sigma) + 1.001 + latent_r = tf.nn.elu(latent_r) + 1.0001 + latent_sigma = tf.nn.elu(latent_sigma) + 1.0001 # normalize scale to sum to 1 latent_scale /= tf.reduce_sum(latent_scale, axis=-1, keepdims=True) From f7509719e0b819d511646d9c84d9a83697aaa39c Mon Sep 17 00:00:00 2001 From: mhuen Date: Mon, 7 Oct 2024 09:19:57 -0500 Subject: [PATCH 15/32] Update file paths in configs --- configs/cascade_11param_noise_ftpv3m.yaml | 13 ----- configs/cascade_7param_noise_ftpv3m.yaml | 58 +++++++++-------------- 2 files changed, 23 insertions(+), 48 deletions(-) diff --git a/configs/cascade_11param_noise_ftpv3m.yaml b/configs/cascade_11param_noise_ftpv3m.yaml index cda06311..1a962c3f 100644 --- a/configs/cascade_11param_noise_ftpv3m.yaml +++ b/configs/cascade_11param_noise_ftpv3m.yaml @@ -253,18 +253,6 @@ data_iterator_settings: { # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30253/cascades/step_3_pass2_get_all_pulses/*/*00000.hdf5', # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30254/cascades/step_3_pass2_get_all_pulses/*/*00000.hdf5', - # # MESE NuGen NuE l5 Benchmark Datasets - # # Spice3.2.1 variable - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99900/egenerator_99900_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # Spice3.2.1 const - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99901/egenerator_99901_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # SpiceLea variable - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99902/egenerator_99902_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # SpiceLea const - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99903/egenerator_99903_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # Spice3.2.1 variable no fourier - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99911/egenerator_99911_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # FTPv3m baseline # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/99915/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', # FTPv3m SnowStorm @@ -384,7 +372,6 @@ data_trafo_settings: { 'float_precision': 'float64', 'norm_constant': !!float 1e-6, 'num_batches': 5000, - # 'model_dir': '/cephfs/users/mhuennefeld/data/egenerator/trafo_models/trafo_model_cascade_11param_noise_ftpv3m', 'model_dir': '/data/user/mhuennefeld/data/egenerator/trafo_models/trafo_model_cascade_11param_noise_ftpv3m', } diff --git a/configs/cascade_7param_noise_ftpv3m.yaml b/configs/cascade_7param_noise_ftpv3m.yaml index 9babdf61..84dd5c17 100644 --- a/configs/cascade_7param_noise_ftpv3m.yaml +++ b/configs/cascade_7param_noise_ftpv3m.yaml @@ -181,11 +181,11 @@ data_iterator_settings: { 'num_add_files': 12, 'num_repetitions': 5, 'input_data': [ - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30250/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30251/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30252/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30250/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30251/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30252/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', ], }, @@ -200,11 +200,11 @@ data_iterator_settings: { 'num_repetitions': 1, 'pick_random_files_forever': False, 'input_data': [ - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30250/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30251/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30252/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30250/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30251/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30252/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', ], }, @@ -221,11 +221,11 @@ data_iterator_settings: { 'pick_random_files_forever': True, 'input_data': [ # Note: The validation data is the same as the training data. Do not pay attention to validation curve! - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30250/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30251/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30252/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30250/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30251/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30252/cascades/step_3_pass2_get_all_pulses/*/*0.hdf5', ], }, @@ -243,31 +243,19 @@ data_iterator_settings: { 'input_data': [ # # validation data - # '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*00000.hdf5', - # '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*00000.hdf5', - - # # MESE NuGen NuE l5 Benchmark Datasets - # # Spice3.2.1 variable - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99900/egenerator_99900_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # Spice3.2.1 const - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99901/egenerator_99901_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # SpiceLea variable - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99902/egenerator_99902_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # SpiceLea const - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99903/egenerator_99903_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', - # # Spice3.2.1 variable no fourier - # # '/net/big-tank/POOL/users/mhuennefeld/data/egenerator/training_data/datasets/99911/egenerator_99911_step_3_pass2_get_pulses_py3_v4_1_0_IC86_pulses/*/00000-00999/*_0000000*.hdf5', + # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30248/cascades/step_3_pass2_get_all_pulses/*/*00000.hdf5', + # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/30249/cascades/step_3_pass2_get_all_pulses/*/*00000.hdf5', # FTPv3m baseline - '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/99915/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/99915/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', # # FTPv3m SnowStorm - # '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/99916/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/99916/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', # # FTPv3m SnowStorm Absorption - # '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/99917/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/99917/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', # # FTPv3m SnowStorm Scattering - # '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/99918/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/99918/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', # # FTPv3m SnowStorm HoleIce - # '/cephfs/projects/ICECUBE/icecube/training_data/egenerator/egenerator-v1.1.0/datasets/99919/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', + # '/data/ana/reconstruction/2018/gnn/training_data/egenerator-v1.1.0/datasets/99919/cascades/step_3_pass2_get_all_pulses/*/*.hdf5', ], }, @@ -357,7 +345,7 @@ data_trafo_settings: { 'float_precision': 'float64', 'norm_constant': !!float 1e-6, 'num_batches': 5000, - 'model_dir': '/cephfs/users/mhuennefeld/data/egenerator/trafo_models/trafo_model_cascade_7param_noise_ftpv3m', + 'model_dir': '/data/user/mhuennefeld/data/egenerator/trafo_models/trafo_model_cascade_7param_noise_ftpv3m', } #---------------------- From cfa478e7133005d0808b795be8b64eb3ef2fd8d6 Mon Sep 17 00:00:00 2001 From: mhuen Date: Mon, 7 Oct 2024 09:32:27 -0500 Subject: [PATCH 16/32] Fix unittest --- egenerator/model/source/base.py | 8 ++++---- test/model/multi_source/test_independent.py | 5 ++++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/egenerator/model/source/base.py b/egenerator/model/source/base.py index 2e467ebd..3b57b4c6 100644 --- a/egenerator/model/source/base.py +++ b/egenerator/model/source/base.py @@ -116,11 +116,11 @@ def epsilon(self): elif "sources" in config: model_precisions = [] - for source in config["sources"]: + for _, base_source in config["sources"].items(): model_precisions.append( - self.sub_components[source].configuration.config["config"][ - "float_precision" - ] + self.sub_components[base_source].configuration.config[ + "config" + ]["float_precision"] ) model_precisions = np.unique(model_precisions) if len(model_precisions) > 1: diff --git a/test/model/multi_source/test_independent.py b/test/model/multi_source/test_independent.py index 049a7e88..e8fc6a4c 100644 --- a/test/model/multi_source/test_independent.py +++ b/test/model/multi_source/test_independent.py @@ -76,7 +76,10 @@ def setUp(self): "cascade_00002_energy", "cascade_00002_time", ] - self.config_cascade = {"cascade_setting": 1337} + self.config_cascade = { + "cascade_setting": 1337, + "float_precision": "float32", + } self.base_sources = { "cascade": self.get_cascade_source( config=self.config_cascade, data_trafo=self.data_trafo From 0a50822e480eca339fad701641b683e8e8146afe Mon Sep 17 00:00:00 2001 From: mhuen Date: Tue, 8 Oct 2024 10:11:48 -0500 Subject: [PATCH 17/32] cascades: Fix float64 --- egenerator/model/source/cascade/default.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/egenerator/model/source/cascade/default.py b/egenerator/model/source/cascade/default.py index f45a7543..99bc49e8 100644 --- a/egenerator/model/source/cascade/default.py +++ b/egenerator/model/source/cascade/default.py @@ -890,7 +890,8 @@ def get_tensors( # ------------------------------------------- pulse_latent_scale = tf.cast( - pulse_latent_scale, dtype=config["float_precision"] + pulse_latent_scale, + dtype="float64", ) # [n_pulses, 1] * [n_pulses, n_models] = [n_pulses, n_models] From 7788bb5339bcc4539b4774142a69369517e3ff1b Mon Sep 17 00:00:00 2001 From: mhuen Date: Tue, 8 Oct 2024 10:27:40 -0500 Subject: [PATCH 18/32] Option to set float precision for pdf cdf evaluation independently --- configs/cascade_11param_noise_ftpv3m.yaml | 1 + configs/cascade_7param_noise_ftpv3m.yaml | 1 + egenerator/model/source/cascade/default.py | 15 +++++++-------- .../model/source/track/inf_tracks_sphere.py | 12 ++++++------ 4 files changed, 15 insertions(+), 14 deletions(-) diff --git a/configs/cascade_11param_noise_ftpv3m.yaml b/configs/cascade_11param_noise_ftpv3m.yaml index 1a962c3f..4efe1777 100644 --- a/configs/cascade_11param_noise_ftpv3m.yaml +++ b/configs/cascade_11param_noise_ftpv3m.yaml @@ -505,6 +505,7 @@ model_settings: { 'charge_distribution_type': 'negative_binomial', 'num_latent_models': 10, 'float_precision': float32, + 'float_precision_pdf_cdf': float64, # Baseline DOM Angular acceptance 'use_constant_baseline_hole_ice': False, diff --git a/configs/cascade_7param_noise_ftpv3m.yaml b/configs/cascade_7param_noise_ftpv3m.yaml index 84dd5c17..dea4db51 100644 --- a/configs/cascade_7param_noise_ftpv3m.yaml +++ b/configs/cascade_7param_noise_ftpv3m.yaml @@ -480,6 +480,7 @@ model_settings: { 'charge_distribution_type': 'negative_binomial', 'num_latent_models': 10, 'float_precision': float32, + 'float_precision_pdf_cdf': float64, # Baseline DOM Angular acceptance 'use_constant_baseline_hole_ice': True, diff --git a/egenerator/model/source/cascade/default.py b/egenerator/model/source/cascade/default.py index 99bc49e8..c8537278 100644 --- a/egenerator/model/source/cascade/default.py +++ b/egenerator/model/source/cascade/default.py @@ -621,14 +621,14 @@ def get_tensors( mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) tw_cdf_stop = basis_functions.tf_asymmetric_gauss_cdf( x=t_exclusions[:, 1], mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) # shape: [n_tw, n_models] @@ -784,7 +784,7 @@ def get_tensors( mu=dom_charges, sigma=dom_charges_sigma, r=dom_charges_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) ), dtype=config["float_precision"], @@ -835,7 +835,7 @@ def get_tensors( x=dom_charges_true, mu=dom_charges, alpha=dom_charges_alpha, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ), dtype=config["float_precision"], ) @@ -890,8 +890,7 @@ def get_tensors( # ------------------------------------------- pulse_latent_scale = tf.cast( - pulse_latent_scale, - dtype="float64", + pulse_latent_scale, dtype=config["float_precision_pdf_cdf"] ) # [n_pulses, 1] * [n_pulses, n_models] = [n_pulses, n_models] @@ -901,7 +900,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) * pulse_latent_scale ) @@ -911,7 +910,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) * pulse_latent_scale ) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index f0ed9e0e..261498b4 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -694,14 +694,14 @@ def get_tensors( mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) tw_cdf_stop = basis_functions.tf_asymmetric_gauss_cdf( x=t_exclusions[:, 1], mu=tw_latent_mu, sigma=tw_latent_sigma, r=tw_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) # shape: [n_tw, n_models] @@ -855,7 +855,7 @@ def get_tensors( mu=dom_charges, sigma=dom_charges_sigma, r=dom_charges_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) ), config["float_precision"], @@ -905,7 +905,7 @@ def get_tensors( x=dom_charges_true, mu=dom_charges, alpha=dom_charges_alpha, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) dom_charges_llh = tf.cast( dom_charges_llh, config["float_precision"] @@ -969,7 +969,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) * pulse_latent_scale ) @@ -979,7 +979,7 @@ def get_tensors( mu=pulse_latent_mu, sigma=pulse_latent_sigma, r=pulse_latent_r, - dtype="float64", + dtype=config["float_precision_pdf_cdf"], ) * pulse_latent_scale ) From 9af9b0784717c7cfc69bb4ee12b50ee4d929eb6c Mon Sep 17 00:00:00 2001 From: mhuen Date: Tue, 8 Oct 2024 11:26:21 -0500 Subject: [PATCH 19/32] Option to set float precision for pdf cdf evaluation independently --- configs/track_sphere_6param_ftpv3m.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/configs/track_sphere_6param_ftpv3m.yaml b/configs/track_sphere_6param_ftpv3m.yaml index 27b3ecd5..03993a0b 100644 --- a/configs/track_sphere_6param_ftpv3m.yaml +++ b/configs/track_sphere_6param_ftpv3m.yaml @@ -500,6 +500,7 @@ model_settings: { 'charge_distribution_type': 'negative_binomial', 'num_latent_models': 10, 'float_precision': float32, + 'float_precision_pdf_cdf': float64, # Baseline DOM Angular acceptance 'use_constant_baseline_hole_ice': True, From 6bec37d093f6563b1603196fa619dd00acb1e051 Mon Sep 17 00:00:00 2001 From: mhuen Date: Wed, 23 Oct 2024 15:48:38 -0500 Subject: [PATCH 20/32] Turn off normalization by total charge --- configs/cascade_11param_noise_ftpv3m.yaml | 2 +- configs/cascade_7param_noise_ftpv3m.yaml | 2 +- configs/track_sphere_6param_ftpv3m.yaml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/configs/cascade_11param_noise_ftpv3m.yaml b/configs/cascade_11param_noise_ftpv3m.yaml index 4efe1777..fa2517e9 100644 --- a/configs/cascade_11param_noise_ftpv3m.yaml +++ b/configs/cascade_11param_noise_ftpv3m.yaml @@ -64,7 +64,7 @@ training_settings: { # Additional keywords to the loss module used for training 'additional_loss_module_kwargs': { - 'normalize_by_total_charge': True, + 'normalize_by_total_charge': False, }, } diff --git a/configs/cascade_7param_noise_ftpv3m.yaml b/configs/cascade_7param_noise_ftpv3m.yaml index dea4db51..fec9bbf6 100644 --- a/configs/cascade_7param_noise_ftpv3m.yaml +++ b/configs/cascade_7param_noise_ftpv3m.yaml @@ -64,7 +64,7 @@ training_settings: { # Additional keywords to the loss module used for training 'additional_loss_module_kwargs': { - 'normalize_by_total_charge': True, + 'normalize_by_total_charge': False, }, } diff --git a/configs/track_sphere_6param_ftpv3m.yaml b/configs/track_sphere_6param_ftpv3m.yaml index 03993a0b..d9485ca3 100644 --- a/configs/track_sphere_6param_ftpv3m.yaml +++ b/configs/track_sphere_6param_ftpv3m.yaml @@ -64,7 +64,7 @@ training_settings: { # Additional keywords to the loss module used for training 'additional_loss_module_kwargs': { - 'normalize_by_total_charge': True, + 'normalize_by_total_charge': False, }, } From a55d24517815056aadba5af7c4cd84f9a3ec7612 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 26 Oct 2024 13:12:32 -0500 Subject: [PATCH 21/32] Add closest approach point info --- .../model/source/track/inf_tracks_sphere.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index 261498b4..8f37b14e 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -241,6 +241,10 @@ def get_tensors( # Locals: # - dx, dy, dz, dist, length_along_track [cherenkov] # - opening angle cherenkov cone and PMT direction + # - opening angle track direction and displacement vector + # - x, y, y [closest approach point to DOM] + # - closest approach distance + # - dx, dy, dz [closest approach point to DOM] # # Note: we will try to keep the number of input features # to a minimum for now. One could test if adding more @@ -302,6 +306,12 @@ def get_tensors( # shape: [n_batch, 86, 60, 1] distance_closest = tf.sqrt(dx_inf**2 + dy_inf**2 + dz_inf**2) + 1e-1 + # normalize displacement vectors + # shape: [n_batch, 86, 60, 1] + dx_inf_normed = dx_inf / distance_closest + dy_inf_normed = dy_inf / distance_closest + dz_inf_normed = dz_inf / distance_closest + # shift to cherenkov emission point # calculate distance on track of cherenkov position # shape: [n_batch, 86, 60, 1] @@ -359,6 +369,13 @@ def get_tensors( ), )[..., tf.newaxis] + # calculate opening angle of track direction and displacement vector + # from the closest approach point to the DOM + opening_angle_closest = angles.get_angle( + tf.stack([dir_x, dir_y, dir_z], axis=-1), + tf.concat([dx_inf, dy_inf, dz_inf], axis=-1), + )[..., tf.newaxis] + # compute t_geometry: time for photon to travel to DOM # Shape: [n_batch, 86, 60, 1] c_ice = 0.22103046286329384 # m/ns @@ -369,6 +386,9 @@ def get_tensors( norm_const = self.data_trafo.data["norm_constant"] # transform distances and lengths in detector + distance_closest_tr = distance_closest / ( + config["sphere_radius"] + norm_const + ) distance_cherenkov_tr = distance_cherenkov / ( config["sphere_radius"] + norm_const ) @@ -376,8 +396,16 @@ def get_tensors( config["sphere_radius"] + norm_const ) + # transform coordinates by approximate size of IceCube + closest_x_tr = closest_x / (500.0 + norm_const) + closest_y_tr = closest_y / (500.0 + norm_const) + closest_z_tr = closest_z / (500.0 + norm_const) + # transform angle opening_angle_traf = opening_angle / (norm_const + np.pi) + opening_angle_closest_traf = opening_angle_closest / ( + norm_const + np.pi + ) # ---------------------- # Collect input features @@ -412,6 +440,14 @@ def get_tensors( distance_cherenkov_tr, length_cherenkov_pos_tr, opening_angle_traf, + opening_angle_closest_traf, + closest_x_tr, + closest_y_tr, + closest_z_tr, + distance_closest_tr, + dx_inf_normed, + dy_inf_normed, + dz_inf_normed, ] if config["add_anisotropy_angle"]: From 82750f54ce21161328c72c9479ad737bdcfaab68 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 26 Oct 2024 13:13:19 -0500 Subject: [PATCH 22/32] Increase default model size: use reco pulses --- configs/track_sphere_6param_ftpv3m.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/configs/track_sphere_6param_ftpv3m.yaml b/configs/track_sphere_6param_ftpv3m.yaml index d9485ca3..b46c3436 100644 --- a/configs/track_sphere_6param_ftpv3m.yaml +++ b/configs/track_sphere_6param_ftpv3m.yaml @@ -380,7 +380,7 @@ data_handler_settings: { # settings for the data module 'data_settings':{ - 'pulse_key': 'MCPulses', + 'pulse_key': 'InIceDSTPulses_masked_doms_only', 'event_id_key': 'LabelsMCTrackSphere', 'dom_exclusions_key': BadDomsList, 'time_exclusions_key': , @@ -488,7 +488,7 @@ model_settings: { 'keep_prob':, 'sphere_radius': 750., 'add_anisotropy_angle': True, - 'add_dom_angular_acceptance': False, + 'add_dom_angular_acceptance': True, 'add_dom_coordinates': False, 'num_local_vars': 0, 'scale_charge': True, @@ -509,7 +509,7 @@ model_settings: { # First convolutions 'filter_size_list' : [[1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1]], - 'num_filters_list' : [25, 100, 100, 100, 100, 42], + 'num_filters_list' : [50, 200, 200, 200, 200, 42], 'method_list' : ['locally_connected', 'convolution', 'convolution', 'convolution', 'convolution', 'convolution', From 3d5635de6c5dcd26574d3179db8fb32619974174 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 26 Oct 2024 13:35:46 -0500 Subject: [PATCH 23/32] Add one earlier seed point for tracks --- egenerator/model/source/track/inf_tracks_sphere.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index 8f37b14e..2902a205 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -659,7 +659,7 @@ def get_tensors( # features. t_seed = ( np.r_[ - [0, 100, 8000, 14000, 4000, 800, 300, 1000, 400, 2000], + [0, -100, 100, 8000, 4000, 800, 300, 1000, 400, 2000], np.random.RandomState(42).uniform(0, 14000, max(1, n_models)), ][:n_models] * t_scale From 98719eb8d587055d4c597916107fad34c60e4c5c Mon Sep 17 00:00:00 2001 From: mhuen Date: Mon, 28 Oct 2024 14:32:14 -0500 Subject: [PATCH 24/32] fix typo --- egenerator/model/source/track/inf_tracks_sphere.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index 2902a205..e3315482 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -372,7 +372,7 @@ def get_tensors( # calculate opening angle of track direction and displacement vector # from the closest approach point to the DOM opening_angle_closest = angles.get_angle( - tf.stack([dir_x, dir_y, dir_z], axis=-1), + tf.concat([dir_x, dir_y, dir_z], axis=-1), tf.concat([dx_inf, dy_inf, dz_inf], axis=-1), )[..., tf.newaxis] From b2c6071ac70c18aaeadf53315a99ed8646ac160b Mon Sep 17 00:00:00 2001 From: mhuen Date: Mon, 28 Oct 2024 14:36:05 -0500 Subject: [PATCH 25/32] fix typo --- egenerator/model/source/track/inf_tracks_sphere.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index e3315482..9fe95025 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -80,7 +80,7 @@ def _build_architecture(self, config, name=None): for i in range(num): parameter_names.append(param_name.format(i)) - num_inputs = 13 + num_snowstorm_params + num_inputs = 21 + num_snowstorm_params if config["add_anisotropy_angle"]: num_inputs += 2 From c6b64663fc5dd6d03c0b636998f6768061faf0e0 Mon Sep 17 00:00:00 2001 From: mhuen Date: Wed, 30 Oct 2024 22:17:33 -0500 Subject: [PATCH 26/32] Prevent division by zero --- egenerator/model/multi_source/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/egenerator/model/multi_source/base.py b/egenerator/model/multi_source/base.py index d76f6d04..21319388 100644 --- a/egenerator/model/multi_source/base.py +++ b/egenerator/model/multi_source/base.py @@ -459,7 +459,7 @@ def get_tensors( # normalize time exclusion sum: divide by total charge at DOM if time_exclusions_exist: dom_cdf_exclusion_sum = tf_helpers.safe_cdf_clip( - dom_cdf_exclusion_sum / dom_charges + dom_cdf_exclusion_sum / (dom_charges + self.epsilon) ) result_tensors["dom_cdf_exclusion_sum"] = dom_cdf_exclusion_sum From b6f3c7c98431838fd9f278b05d5ac24e13a329c5 Mon Sep 17 00:00:00 2001 From: mhuen Date: Wed, 30 Oct 2024 22:18:12 -0500 Subject: [PATCH 27/32] Add buffer for time window calculation --- configs/cascade_11param_noise_ftpv3m.yaml | 1 + configs/cascade_7param_noise_ftpv3m.yaml | 1 + configs/track_sphere_6param_ftpv3m.yaml | 1 + egenerator/data/modules/data/pulse_data.py | 22 ++++++++++++++++++++++ egenerator/manager/source.py | 4 +++- 5 files changed, 28 insertions(+), 1 deletion(-) diff --git a/configs/cascade_11param_noise_ftpv3m.yaml b/configs/cascade_11param_noise_ftpv3m.yaml index fa2517e9..c4c3b88b 100644 --- a/configs/cascade_11param_noise_ftpv3m.yaml +++ b/configs/cascade_11param_noise_ftpv3m.yaml @@ -400,6 +400,7 @@ data_handler_settings: { 'float_precision': 'float32', 'add_charge_quantiles': False, 'discard_pulses_from_excluded_doms': False, + 'time_window_buffer': 100., }, # -------------------- diff --git a/configs/cascade_7param_noise_ftpv3m.yaml b/configs/cascade_7param_noise_ftpv3m.yaml index fec9bbf6..86f4a0e2 100644 --- a/configs/cascade_7param_noise_ftpv3m.yaml +++ b/configs/cascade_7param_noise_ftpv3m.yaml @@ -373,6 +373,7 @@ data_handler_settings: { 'float_precision': 'float32', 'add_charge_quantiles': False, 'discard_pulses_from_excluded_doms': False, + 'time_window_buffer': 100., }, # -------------------- diff --git a/configs/track_sphere_6param_ftpv3m.yaml b/configs/track_sphere_6param_ftpv3m.yaml index b46c3436..0855a881 100644 --- a/configs/track_sphere_6param_ftpv3m.yaml +++ b/configs/track_sphere_6param_ftpv3m.yaml @@ -387,6 +387,7 @@ data_handler_settings: { 'float_precision': 'float32', 'add_charge_quantiles': False, 'discard_pulses_from_excluded_doms': False, + 'time_window_buffer': 100., }, # -------------------- diff --git a/egenerator/data/modules/data/pulse_data.py b/egenerator/data/modules/data/pulse_data.py index 225385c9..91692ba0 100644 --- a/egenerator/data/modules/data/pulse_data.py +++ b/egenerator/data/modules/data/pulse_data.py @@ -44,6 +44,7 @@ def _configure( float_precision, add_charge_quantiles, discard_pulses_from_excluded_doms, + time_window_buffer, ): """Configure Module Class This is an abstract method and must be implemented by derived class. @@ -75,6 +76,11 @@ def _configure( discard_pulses_from_excluded_doms : bool, optional If True, pulses on excluded DOMs are discarded. The pulses are discarded after the charge at the DOM is collected. + time_window_buffer : float + The additional buffer time to add to the time window + on either side. The time window is defined as + [min(pulse_times) - buffer, max(pulse_times) + buffer]. + Returns ------- @@ -120,6 +126,9 @@ def _configure( msg = "Pulse key type: {!r} != str" raise TypeError(msg.format(type(pulse_key))) + if time_window_buffer < 0: + raise ValueError("Time window buffer must be >= 0") + time_exclusions_exist = time_exclusions_key is not None dom_exclusions_exist = dom_exclusions_key is not None @@ -214,6 +223,7 @@ def _configure( config_data=config_data, float_precision=float_precision, add_charge_quantiles=add_charge_quantiles, + time_window_buffer=time_window_buffer, ), mutable_settings=dict( pulse_key=pulse_key, @@ -392,6 +402,12 @@ def get_data_from_hdf(self, file, *args, **kwargs): if row.time < x_time_window[index, 0]: x_time_window[index, 0] = row.time + # increase time window by buffer + buffer = self.configuration.config["time_window_buffer"] + if buffer > 0: + x_time_window[:, 0] -= buffer + x_time_window[:, 1] += buffer + # convert cumulative charge to fraction of total charge, e.g. quantile if add_charge_quantiles: @@ -624,6 +640,12 @@ def get_data_from_frame(self, frame, *args, **kwargs): if pulse.time < x_time_window[index, 0]: x_time_window[index, 0] = pulse.time + # increase time window by buffer + buffer = self.configuration.config["time_window_buffer"] + if buffer > 0: + x_time_window[:, 0] -= buffer + x_time_window[:, 1] += buffer + x_pulses = np.array(x_pulses, dtype=self.data["np_float_precision"]) x_pulses_ids = np.array(x_pulses_ids, dtype=np.int32) diff --git a/egenerator/manager/source.py b/egenerator/manager/source.py index b8d0e838..3464cd11 100644 --- a/egenerator/manager/source.py +++ b/egenerator/manager/source.py @@ -635,7 +635,9 @@ def model_tensors_function( x_pulses_ids=tf.convert_to_tensor([[0, 0, 0, 0]]), x_dom_exclusions=tf.ones([0, 86, 60, 1], dtype=tf.bool), x_dom_charge=tf.ones([0, 86, 60, 1], dtype=param_dtype), - x_time_window=tf.convert_to_tensor([[9000, 9001]], dtype=tw_dtype), + x_time_window=tf.convert_to_tensor( + [[6000, 15000]], dtype=tw_dtype + ), x_time_exclusions=tf.convert_to_tensor( [[0, 0]], dtype=t_exclusions_dtype ), From f1a64e791196185bdb825bb2e94d7d994eb1530b Mon Sep 17 00:00:00 2001 From: mhuen Date: Wed, 30 Oct 2024 22:58:46 -0500 Subject: [PATCH 28/32] Add buffer for time window calculation: handle no pulse case --- egenerator/data/modules/data/pulse_data.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/egenerator/data/modules/data/pulse_data.py b/egenerator/data/modules/data/pulse_data.py index 91692ba0..48f8e025 100644 --- a/egenerator/data/modules/data/pulse_data.py +++ b/egenerator/data/modules/data/pulse_data.py @@ -404,6 +404,14 @@ def get_data_from_hdf(self, file, *args, **kwargs): # increase time window by buffer buffer = self.configuration.config["time_window_buffer"] + + # non-finite values should only happen if there are no pulses + mask_infinite1 = ~np.isfinite(x_time_window[:, 0]) + mask_infinite2 = ~np.isfinite(x_time_window[:, 1]) + assert np.all(mask_infinite1 == mask_infinite2) + x_time_window[mask_infinite1, 0] = 0 + x_time_window[mask_infinite1, 1] = 0 + if buffer > 0: x_time_window[:, 0] -= buffer x_time_window[:, 1] += buffer @@ -642,6 +650,14 @@ def get_data_from_frame(self, frame, *args, **kwargs): # increase time window by buffer buffer = self.configuration.config["time_window_buffer"] + + # non-finite values should only happen if there are no pulses + mask_infinite1 = ~np.isfinite(x_time_window[:, 0]) + mask_infinite2 = ~np.isfinite(x_time_window[:, 1]) + assert np.all(mask_infinite1 == mask_infinite2) + x_time_window[mask_infinite1, 0] = 0 + x_time_window[mask_infinite1, 1] = 0 + if buffer > 0: x_time_window[:, 0] -= buffer x_time_window[:, 1] += buffer From 50ac3b1b0e6aca18be6c1d23c3a520f7c47746e4 Mon Sep 17 00:00:00 2001 From: mhuen Date: Thu, 31 Oct 2024 08:30:09 -0500 Subject: [PATCH 29/32] Typo --- egenerator/data/modules/data/pulse_data.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/egenerator/data/modules/data/pulse_data.py b/egenerator/data/modules/data/pulse_data.py index 48f8e025..a319fbff 100644 --- a/egenerator/data/modules/data/pulse_data.py +++ b/egenerator/data/modules/data/pulse_data.py @@ -410,7 +410,7 @@ def get_data_from_hdf(self, file, *args, **kwargs): mask_infinite2 = ~np.isfinite(x_time_window[:, 1]) assert np.all(mask_infinite1 == mask_infinite2) x_time_window[mask_infinite1, 0] = 0 - x_time_window[mask_infinite1, 1] = 0 + x_time_window[mask_infinite2, 1] = 0 if buffer > 0: x_time_window[:, 0] -= buffer @@ -656,7 +656,7 @@ def get_data_from_frame(self, frame, *args, **kwargs): mask_infinite2 = ~np.isfinite(x_time_window[:, 1]) assert np.all(mask_infinite1 == mask_infinite2) x_time_window[mask_infinite1, 0] = 0 - x_time_window[mask_infinite1, 1] = 0 + x_time_window[mask_infinite2, 1] = 0 if buffer > 0: x_time_window[:, 0] -= buffer From 2cf346128989219979019cbab12ab9ec4c3db043 Mon Sep 17 00:00:00 2001 From: mhuen Date: Thu, 31 Oct 2024 13:17:41 -0500 Subject: [PATCH 30/32] Reco speedup: fix data_batch and seed_array for minimizer function --- egenerator/manager/source.py | 136 +++++++++++++++++------------------ 1 file changed, 68 insertions(+), 68 deletions(-) diff --git a/egenerator/manager/source.py b/egenerator/manager/source.py index 3464cd11..0c83b71c 100644 --- a/egenerator/manager/source.py +++ b/egenerator/manager/source.py @@ -796,11 +796,44 @@ def reconstruct_events( msg.format(param_tensor.shape[1], len(fit_parameter_list)) ) + # transform seed if minimization is performed in trafo space + if isinstance(seed, str): + seed_index = self.data_handler.tensors.get_index(seed) + seed_array = data_batch[seed_index] + else: + seed_array = seed + if minimize_in_trafo_space: + + # transform bounds if provided + if "bounds" in kwargs: + bounds = self.data_trafo.transform( + data=np.array(kwargs["bounds"]).T, + tensor_name=parameter_tensor_name, + ).T + for i, bound in enumerate(bounds): + for j in range(2): + if not np.isfinite(bound[j]): + bounds[i, j] = None + kwargs["bounds"] = bounds + + seed_array_trafo = self.data_trafo.transform( + data=seed_array, tensor_name=parameter_tensor_name + ) + else: + seed_array_trafo = seed_array + + # get seed parameters + if np.all(fit_parameter_list): + x0 = seed_array_trafo + else: + # get seed parameters + x0 = seed_array_trafo[:, fit_parameter_list] + # define helper function - def func(x, data_batch, seed): + def func(x): # reshape and convert to proper x = np.reshape(x, param_shape).astype(param_dtype) - seed = np.reshape(seed, param_shape_full).astype(param_dtype) + seed = np.reshape(seed_array, param_shape_full).astype(param_dtype) loss, grad = loss_and_gradients_function(x, data_batch, seed=seed) loss = loss.numpy().astype("float64") grad = grad.numpy().astype("float64") @@ -810,10 +843,12 @@ def func(x, data_batch, seed): if hessian_function is not None: - def get_hessian(x, data_batch, seed): + def get_hessian(x): # reshape and convert to tensor x = np.reshape(x, param_shape).astype(param_dtype) - seed = np.reshape(seed, param_shape_full).astype(param_dtype) + seed = np.reshape(seed_array, param_shape_full).astype( + param_dtype + ) hessian = hessian_function(x, data_batch, seed=seed) hessian = hessian.numpy().astype("float64") return hessian @@ -844,47 +879,9 @@ def get_hessian(x, data_batch, seed): # kwargs['callback'] = tolerance_func - # transform seed if minimization is performed in trafo space - if isinstance(seed, str): - seed_index = self.data_handler.tensors.get_index(seed) - seed_array = data_batch[seed_index] - else: - seed_array = seed - if minimize_in_trafo_space: - - # transform bounds if provided - if "bounds" in kwargs: - bounds = self.data_trafo.transform( - data=np.array(kwargs["bounds"]).T, - tensor_name=parameter_tensor_name, - ).T - for i, bound in enumerate(bounds): - for j in range(2): - if not np.isfinite(bound[j]): - bounds[i, j] = None - kwargs["bounds"] = bounds - - seed_array_trafo = self.data_trafo.transform( - data=seed_array, tensor_name=parameter_tensor_name - ) - else: - seed_array_trafo = seed_array - - # get seed parameters - if np.all(fit_parameter_list): - x0 = seed_array_trafo - else: - # get seed parameters - x0 = seed_array_trafo[:, fit_parameter_list] - x0_flat = np.reshape(x0, [-1]) result = optimize.minimize( - fun=func, - x0=x0_flat, - jac=jac, - method=method, - args=(data_batch, seed_array), - **kwargs + fun=func, x0=x0_flat, jac=jac, method=method, **kwargs ) best_fit = np.reshape(result.x, param_shape) @@ -1121,11 +1118,31 @@ def scipy_global_reconstruct_events( minimizer_kwargs["jac"] = jac options["jac"] = jac + # get seed tensor + if isinstance(seed, str): + seed_index = self.data_handler.tensors.get_index(seed) + seed_array = data_batch[seed_index] + else: + seed_array = seed + + # transform seed if minimization is performed in trafo space + if minimize_in_trafo_space: + seed_array_trafo = self.data_trafo.transform( + data=seed_array, tensor_name=parameter_tensor_name + ) + else: + seed_array_trafo = seed_array + + # For now: add +- 1 in trafo space + # ToDo: allow to pass proper boundaries and uncertainties + assert minimize_in_trafo_space, "currently only for trafo space" + bounds = np.concatenate((seed_array_trafo - 1, seed_array_trafo + 1)).T + # define helper function - def func(x, data_batch, seed): + def func(x): # reshape and convert to tensor x = np.reshape(x, param_shape).astype(param_dtype) - seed = np.reshape(seed, param_shape_full).astype(param_dtype) + seed = np.reshape(seed_array, param_shape_full).astype(param_dtype) loss, grad = loss_and_gradients_function(x, data_batch, seed=seed) loss = loss.numpy().astype("float64") grad = grad.numpy().astype("float64") @@ -1135,10 +1152,12 @@ def func(x, data_batch, seed): if hessian_function is not None: - def get_hessian(x, data_batch, seed): + def get_hessian(x): # reshape and convert to tensor x = np.reshape(x, param_shape).astype(param_dtype) - seed = np.reshape(seed, param_shape_full).astype(param_dtype) + seed = np.reshape(seed_array, param_shape_full).astype( + param_dtype + ) hessian = hessian_function(x, data_batch, seed=seed) hessian = hessian.numpy().astype("float64") return hessian @@ -1146,26 +1165,6 @@ def get_hessian(x, data_batch, seed): minimizer_kwargs["hess"] = get_hessian options["hess"] = get_hessian - # get seed tensor - if isinstance(seed, str): - seed_index = self.data_handler.tensors.get_index(seed) - seed_array = data_batch[seed_index] - else: - seed_array = seed - - # transform seed if minimization is performed in trafo space - if minimize_in_trafo_space: - seed_array_trafo = self.data_trafo.transform( - data=seed_array, tensor_name=parameter_tensor_name - ) - else: - seed_array_trafo = seed_array - - # For now: add +- 1 in trafo space - # ToDo: allow to pass proper boundaries and uncertainties - assert minimize_in_trafo_space, "currently only for trafo space" - bounds = np.concatenate((seed_array_trafo - 1, seed_array_trafo + 1)).T - def callback(xk): print(xk) @@ -1175,7 +1174,6 @@ def callback(xk): options=options, minimizer_kwargs=minimizer_kwargs, callback=callback, - args=(data_batch, seed_array), **kwargs ) @@ -1282,6 +1280,8 @@ def const_loss_and_gradients_function(x): # convert to tensors loss, grad = loss_and_gradients_function(x, data_batch, seed_array) loss = tf.reshape(loss, [1]) + loss = tf.cast(loss, param_tensor.dtype_tf) + grad = tf.cast(grad, param_tensor.dtype_tf) return loss, grad if hessian_function is not None: From 7546fc0fa1c143b4ac87ebf6f035999f53c768ee Mon Sep 17 00:00:00 2001 From: mhuen Date: Thu, 31 Oct 2024 14:07:44 -0500 Subject: [PATCH 31/32] Speedup: use loss only for jac=False --- .../reconstruction/modules/reconstruction.py | 8 +++- egenerator/manager/source.py | 41 ++++++++++++++----- 2 files changed, 38 insertions(+), 11 deletions(-) diff --git a/egenerator/manager/reconstruction/modules/reconstruction.py b/egenerator/manager/reconstruction/modules/reconstruction.py index a9aee4fe..cce9e12f 100644 --- a/egenerator/manager/reconstruction/modules/reconstruction.py +++ b/egenerator/manager/reconstruction/modules/reconstruction.py @@ -166,11 +166,17 @@ def __init__( # choose reconstruction method depending on the optimizer interface if reco_optimizer_interface.lower() == "scipy": + # choose function according to jac (default: True) + scipy_loss_function = loss_and_gradients_function + if "jac" in scipy_optimizer_settings: + if not scipy_optimizer_settings["jac"]: + scipy_loss_function = self.parameter_loss_function + def reconstruction_method(data_batch, seed_tensor): return manager.reconstruct_events( data_batch, loss_module, - loss_and_gradients_function=loss_and_gradients_function, + loss_and_gradients_function=scipy_loss_function, fit_parameter_list=fit_parameter_list, minimize_in_trafo_space=minimize_in_trafo_space, seed=seed_tensor, diff --git a/egenerator/manager/source.py b/egenerator/manager/source.py index 0c83b71c..ec798811 100644 --- a/egenerator/manager/source.py +++ b/egenerator/manager/source.py @@ -743,7 +743,10 @@ def reconstruct_events( method. loss_and_gradients_function : tf.function The tensorflow function: - f(parameters, data_batch, seed_tensor) -> loss, gradients + if jac=True: + f(parameters, data_batch, seed_tensor) -> loss, gradients + else: + f(parameters, data_batch, seed_tensor) -> loss Note: it is imperative that this function uses the same settings for trafo space! fit_parameter_list : bool or list of bool, optional @@ -830,16 +833,34 @@ def reconstruct_events( x0 = seed_array_trafo[:, fit_parameter_list] # define helper function - def func(x): - # reshape and convert to proper - x = np.reshape(x, param_shape).astype(param_dtype) - seed = np.reshape(seed_array, param_shape_full).astype(param_dtype) - loss, grad = loss_and_gradients_function(x, data_batch, seed=seed) - loss = loss.numpy().astype("float64") - grad = grad.numpy().astype("float64") + if jac: - grad_flat = np.reshape(grad, [-1]) - return loss, grad_flat + def func(x): + # reshape and convert to proper + x = np.reshape(x, param_shape).astype(param_dtype) + seed = np.reshape(seed_array, param_shape_full).astype( + param_dtype + ) + loss, grad = loss_and_gradients_function( + x, data_batch, seed=seed + ) + loss = loss.numpy().astype("float64") + grad = grad.numpy().astype("float64") + + grad_flat = np.reshape(grad, [-1]) + return loss, grad_flat + + else: + + def func(x): + # reshape and convert to proper + x = np.reshape(x, param_shape).astype(param_dtype) + seed = np.reshape(seed_array, param_shape_full).astype( + param_dtype + ) + loss = loss_and_gradients_function(x, data_batch, seed=seed) + loss = loss.numpy().astype("float64") + return loss if hessian_function is not None: From 7a829535f2c8e49727d106e58de85f890940f939 Mon Sep 17 00:00:00 2001 From: mhuen Date: Sat, 2 Nov 2024 23:03:57 -0500 Subject: [PATCH 32/32] Bugfix: fixed dt_geometry due to wrong entry_point --- egenerator/model/source/track/inf_tracks_sphere.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/egenerator/model/source/track/inf_tracks_sphere.py b/egenerator/model/source/track/inf_tracks_sphere.py index 9fe95025..a0382292 100644 --- a/egenerator/model/source/track/inf_tracks_sphere.py +++ b/egenerator/model/source/track/inf_tracks_sphere.py @@ -263,9 +263,9 @@ def get_tensors( # calculate normalized vector to entry point # shape: [n_batch, 1, 1, 1] - e_dir_x = -tf.sin(e_zenith) * tf.cos(e_azimuth) - e_dir_y = -tf.sin(e_zenith) * tf.sin(e_azimuth) - e_dir_z = -tf.cos(e_zenith) + e_dir_x = tf.sin(e_zenith) * tf.cos(e_azimuth) + e_dir_y = tf.sin(e_zenith) * tf.sin(e_azimuth) + e_dir_z = tf.cos(e_zenith) # calculate entry position on a sphere of the given radius # shape: [n_batch, 1, 1, 1]