add tf example-008 (full db, USU on abs, but not on shape)

examples/tensorflow/example-008/01_model_preparation.py

@@ -0,0 +1,158 @@
+import sys
+sys.path.append('../../..')
+import pandas as pd
+from scipy.sparse import block_diag
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+from gmapy.data_management.object_utils import (
+    save_objects
+)
+from gmapy.data_management.uncfuns import (
+    create_experimental_covmat,
+    create_datablock_covmat_list,
+    create_prior_covmat
+)
+from gmapy.mappings.tf.compound_map_tf \
+    import CompoundMap as CompoundMapTF
+from gmapy.data_management.database_IO import read_gma_database
+from gmapy.data_management.tablefuns import (
+    create_prior_table,
+    create_experiment_table,
+)
+from gmapy.mappings.priortools import (
+    attach_shape_prior,
+    initialize_shape_prior,
+    remove_dummy_datasets
+)
+from gmapy.mappings.tf.restricted_map import RestrictedMap
+from gmapy.mappings.tf.energy_dependent_absolute_usu_map_tf import (
+    EnergyDependentAbsoluteUSUMap,
+    create_endep_abs_usu_df
+)
+from gmapy.tf_uq.custom_distributions import (
+    MultivariateNormal,
+    MultivariateNormalLikelihoodWithCovParams,
+    DistributionForParameterSubset,
+    UnnormalizedDistributionProduct
+)
+
+tfd = tfp.distributions
+tfb = tfp.bijectors
+
+# retrieve prior estimates and covariances from the database
+db_path = '../../../tests/testdata/data_and_sacs.json'
+db = read_gma_database(db_path)
+remove_dummy_datasets(db['datablock_list'])
+
+priortable = create_prior_table(db['prior_list'])
+priorcov = create_prior_covmat(db['prior_list'])
+
+# prepare experimental quantities
+exptable = create_experiment_table(db['datablock_list'])
+expcov = create_experimental_covmat(db['datablock_list'])
+exptable['UNC'] = np.sqrt(expcov.diagonal())
+
+# For testing: perturb one dataset and see if we get a change
+# exptable.loc[exptable.NODE == 'exp_1025', 'DATA'] *= 1.5
+
+# initialize the normalization errors
+priortable, priorcov = attach_shape_prior((priortable, exptable), covmat=priorcov, raise_if_exists=False)
+compmap = CompoundMapTF((priortable, exptable), reduce=True)
+initialize_shape_prior((priortable, exptable), compmap)
+
+# some convenient shortcuts
+priorvals = priortable.PRIOR.to_numpy()
+expvals = exptable.DATA.to_numpy()
+
+# speed up the pdf log_prob calculations exploiting the block diagonal structure
+expcov_list, idcs_tuples = create_datablock_covmat_list(db['datablock_list'], relative=True)
+expchol_list = [tf.linalg.cholesky(x.toarray()) for x in expcov_list]
+expchol_op_list = [tf.linalg.LinearOperatorLowerTriangular(
+        x, is_non_singular=True, is_square=True
+    ) for x in expchol_list]
+
+# generate a restricted mapping blending out the fixed parameters
+is_adj = priorcov.diagonal() != 0.
+adj_idcs = np.where(is_adj)[0]
+fixed_idcs = np.where(~is_adj)[0]
+restrimap = RestrictedMap(
+    len(priorvals), compmap.propagate, compmap.jacobian,
+    fixed_params=priorvals[fixed_idcs], fixed_params_idcs=fixed_idcs
+)
+propfun = tf.function(restrimap.propagate)
+jacfun = tf.function(restrimap.jacobian)
+
+# generate the experimental covariance matrix
+expcov_chol = tf.linalg.LinearOperatorBlockDiag(
+    expchol_op_list, is_non_singular=True, is_square=True)
+expcov_linop = tf.linalg.LinearOperatorComposition(
+    [expcov_chol, expcov_chol.adjoint()],
+    is_self_adjoint=True, is_positive_definite=True
+)
+# generate the USU mapping
+usu_df = create_endep_abs_usu_df(
+        exptable, ('MT:1-R1:8', 'MT:1-R1:9', 'MT:3-R1:9-R2:8'),
+        (0., 1., 6., 12., 17.5), (1e-2, 1e-2, 1e-2, 1e-2, 1e-2)
+)
+usu_map = EnergyDependentAbsoluteUSUMap((usu_df, exptable), reduce=True)
+usu_jac = tf.sparse.to_dense(usu_map.jacobian(usu_df.PRIOR.to_numpy()))
+
+
+def create_like_cov_fun(usu_df, expcov_linop, Smat):
+
+    def map_uncertainties(u):
+        ids = np.zeros((len(usu_df),), dtype=np.int32)
+        for index, row in red_usu_df.iterrows(): 
+            reac = row.REAC
+            energy = row.ENERGY
+            cur_idcs = usu_df.index[
+                (usu_df.REAC == reac) & (usu_df.ENERGY == energy)
+            ].to_numpy()
+            ids[cur_idcs] = index
+        # scatter the uncertainties to the appropriate places
+        tf_ids = tf.constant(ids, dtype=tf.int32)
+        uncs = tf.nn.embedding_lookup(u, tf_ids)
+        return uncs
+
+    def like_cov_fun(u):
+        uncs = map_uncertainties(u)
+        # covop = tf.linalg.LinearOperatorLowRankUpdate(
+        covop = tf.linalg.LinearOperatorLowRankUpdate(
+            expcov_linop, Smat, tf.square(uncs) + 1e-7,
+            is_self_adjoint=True, is_positive_definite=True,
+            is_diag_update_positive=True
+        )
+        return covop
+    red_usu_df = usu_df[['REAC', 'ENERGY']].drop_duplicates()
+    red_usu_df.sort_values(
+        ['REAC', 'ENERGY'], ascending=True, ignore_index=True, inplace=True
+    )
+    return like_cov_fun, red_usu_df
+
+
+like_cov_fun, red_usu_df = create_like_cov_fun(usu_df, expcov_linop, usu_jac)
+num_covpars = len(red_usu_df)
+
+# generate the prior distribution
+is_adj_constr = is_adj & np.isfinite(priorcov.diagonal())
+is_adj_constr_idcs = np.where(is_adj_constr)[0]
+priorcov_chol = tf.linalg.LinearOperatorDiag(np.sqrt(priorcov.diagonal()[is_adj_constr]))
+prior_red = MultivariateNormal(priorvals[is_adj_constr], priorcov_chol)
+prior_red.log_prob_hessian(priorvals[is_adj_constr])
+prior = DistributionForParameterSubset(
+    prior_red, len(adj_idcs) + num_covpars, is_adj_constr_idcs
+)
+
+# generate the likelihood
+likelihood = MultivariateNormalLikelihoodWithCovParams(
+    len(adj_idcs), num_covpars, propfun, jacfun, expvals, like_cov_fun,
+    approximate_hessian=True, relative=True
+)
+
+# combine prior and likelihood into posterior
+post = UnnormalizedDistributionProduct([prior, likelihood])
+
+save_objects('output/01_model_preparation_output.pkl', locals(),
+             'post', 'likelihood', 'priorvals', 'is_adj', 'usu_df', 'red_usu_df',
+             'num_covpars', 'priortable', 'exptable', 'expcov', 'like_cov_fun', 'compmap', 'restrimap')

examples/tensorflow/example-008/02_parameter_optimization.py

@@ -0,0 +1,57 @@
+import sys
+sys.path.append('../../..')
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import tensorflow_probability as tfp
+from gmapy.tf_uq.auxiliary import (
+    make_positive_definite, invert_symmetric_matrix
+)
+from gmapy.data_management.object_utils import (
+    load_objects, save_objects
+)
+
+post, likelihood, priorvals, is_adj, usu_df, red_usu_df, num_covpars = \
+    load_objects('output/01_model_preparation_output.pkl',
+                 'post', 'likelihood', 'priorvals', 'is_adj',
+                 'usu_df', 'red_usu_df', 'num_covpars')
+
+# speed it up!
+neg_log_prob_and_gradient = tf.function(post.neg_log_prob_and_gradient)
+
+covpars = np.full(num_covpars, 1.)
+refvals = likelihood.combine_pars(priorvals[is_adj], covpars)
+max_outer_iters = 20
+# max_inner_iters = np.round(np.linspace(50, 500, max_outer_iters)).astype(np.int32)
+max_inner_iters = np.full(max_outer_iters, 500, dtype=np.int32)
+outer_iter = 0
+converged = False
+while not converged and outer_iter < max_outer_iters:
+    max_inner_iter = tf.constant(max_inner_iters[outer_iter], dtype=tf.int32)
+    outer_iter += 1
+    print(f'#  outer iteration {outer_iter}')
+    print(f'-- running inner iteration with {max_inner_iter} iterations')
+    # obtain an approximation of the posterior covariance matrix to aid optimization
+    neg_log_post_hessian = post.neg_log_prob_hessian(refvals)
+    fixed_neg_log_post_hessian = make_positive_definite(neg_log_post_hessian, 1e-4)
+    inv_neg_log_post_hessian = invert_symmetric_matrix(fixed_neg_log_post_hessian)
+    fixed_inv_neg_log_post_hessian = make_positive_definite(inv_neg_log_post_hessian, 1e-4)
+    # find peak of posterior distribution (to use it as a starting value of MCMC)
+    optres = tfp.optimizer.bfgs_minimize(
+        neg_log_prob_and_gradient, initial_position=refvals,
+        initial_inverse_hessian_estimate=fixed_inv_neg_log_post_hessian,
+        max_iterations=max_inner_iter
+    )
+    converged = optres.converged.numpy()
+    refvals = optres.position
+    print(optres)
+    print('USU parameters')
+    red_usu_df['USU'] = optres.position.numpy()[-num_covpars:]
+    print(red_usu_df)
+
+
+# save the optimized parameters
+params, covpars = likelihood.split_pars(refvals)
+
+save_objects('output/02_parameter_optimization_output.pkl', locals(),
+             'optres', 'params', 'covpars', 'usu_df', 'red_usu_df')

examples/tensorflow/example-008/03_mcmc_sampling.py

@@ -0,0 +1,112 @@
+import sys
+sys.path.append('../../..')
+import time
+import gc
+import tensorflow as tf
+import tensorflow_probability as tfp
+from gmapy.data_management.object_utils import (
+    load_objects, save_objects
+)
+from gmapy.tf_uq.auxiliary import (
+    make_positive_definite,
+    invert_symmetric_matrix
+)
+# NOTE: the numpy import is necessary because
+#       a method in the post instance loaded via dill
+#       depends on it
+import numpy as np
+import pandas as pd
+tfb = tfp.bijectors
+
+
+post, likelihood = load_objects(
+    'output/01_model_preparation_output.pkl',
+    'post', 'likelihood'
+)
+optres, = load_objects(
+    'output/02_parameter_optimization_output.pkl',
+    'optres'
+)
+
+# define essential input quantities for MCMC
+optvals = optres.position
+neg_log_post_hessian = post.neg_log_prob_hessian(optvals)
+neg_log_post_hessian = make_positive_definite(neg_log_post_hessian, 1e-10)
+inv_neg_log_post_hessian = invert_symmetric_matrix(neg_log_post_hessian)
+inv_neg_log_post_hessian = make_positive_definite(inv_neg_log_post_hessian, 1e-10)
+
+postcov_chol = tf.linalg.cholesky(inv_neg_log_post_hessian)
+target_log_prob = likelihood.log_prob
+del neg_log_post_hessian
+del inv_neg_log_post_hessian
+gc.collect()
+
+# try the Hamilton Monte Carlo approach
+# num_results = int(1e3)
+# num_burnin_steps = int(2e3)
+num_results = int(20e3)
+num_burnin_steps = int(2e3)
+
+hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
+    target_log_prob_fn=target_log_prob,
+    step_size=0.001,
+    num_leapfrog_steps=10
+)
+
+precond_bijector = tfb.ScaleMatvecTriL(
+    scale_tril=postcov_chol
+)
+
+trafo_hmc_kernel = tfp.mcmc.TransformedTransitionKernel(
+    hmc_kernel, precond_bijector
+)
+
+
+# NOTE: This class is defined so that the one_step method can be
+# decorated with tf.function to speed up the MCMC sampling.
+# Decorating run_chain below with tf.function leads to
+# main memory usage exceeding 30 GB and the process gets killed.
+class MySimpleStepSizeAdaptation(tfp.mcmc.SimpleStepSizeAdaptation):
+
+    # NOTE: Decoration with tf.function triggers a warning
+    #       that the function cannot be converted and will run
+    #       as is. However, the observed performance indicates
+    #       that the warning message is false and the
+    #       function gets compiled.
+    @tf.function
+    def one_step(self, *state_and_results, **one_step_kwargs):
+        return super().one_step(*state_and_results, **one_step_kwargs)
+
+
+# adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
+adaptive_hmc = MySimpleStepSizeAdaptation(
+    inner_kernel=trafo_hmc_kernel,
+    num_adaptation_steps=int(num_burnin_steps * 0.8)
+)
+
+
+def trace_everything(states, previous_kernel_results):
+    return previous_kernel_results
+
+
+# @tf.function
+def run_chain():
+    samples, tracing_info = tfp.mcmc.sample_chain(
+        num_results=num_results,
+        num_burnin_steps=num_burnin_steps,
+        current_state=optvals,
+        kernel=adaptive_hmc,
+        trace_fn=trace_everything
+    )
+    return samples, tracing_info
+
+
+s1 = time.time()
+chain, tracing_info = run_chain()
+s2 = time.time()
+print(f's2-s1: {s2-s1}')
+
+save_objects(
+    'output/03_mcmc_sampling_output.pkl', locals(),
+    'chain', 'tracing_info'
+)