Source code for n3fit.performfit
"""
Fit action controller
"""
# Backend-independent imports
import copy
import logging
import numpy as np
import n3fit.checks
from n3fit.vpinterface import N3PDF
log = logging.getLogger(__name__)
# Action to be called by validphys
# All information defining the NN should come here in the "parameters" dict
[docs]@n3fit.checks.can_run_multiple_replicas
@n3fit.checks.check_multireplica_qed
@n3fit.checks.check_fiatlux_pdfs_id
@n3fit.checks.check_polarized_configs
def performfit(
*,
experiments_data,
n3fit_checks_action, # wrapper for all checks
replicas, # checks specific to performfit
replicas_nnseed_fitting_data_dict,
posdatasets_fitting_pos_dict,
integdatasets_fitting_integ_dict,
theoryid,
fiatlux,
basis,
fitbasis,
positivity_bound,
sum_rules=True,
parameters,
replica_path,
output_path,
save=None,
load=None,
hyperscanner=None,
hyperopt=None,
kfold_parameters,
tensorboard=None,
debug=False,
maxcores=None,
double_precision=False,
parallel_models=False,
):
"""
This action will (upon having read a validcard) process a full PDF fit
for a set of replicas.
The input to this function is provided by validphys
and/or defined in the runcards or commandline arguments.
This controller is provided with:
1. Seeds generated using the replica number and the seeds defined in the runcard.
2. Loaded datasets with replicas generated.
2.1 Loaded positivity/integrability sets.
The workflow of this controller is as follows:
1. Generate a ModelTrainer object holding information to create the NN and perform a fit
(at this point no NN object has been generated)
1.1 (if hyperopt) generates the hyperopt scanning dictionary
taking as a base the fitting dictionary and the runcard's hyperscanner dictionary
2. Pass the dictionary of parameters to ModelTrainer
for the NN to be generated and the fit performed
2.1 (if hyperopt) Loop over point 4 for `hyperopt` number of times
3. Once the fit is finished, output the PDF grid and accompanying files
Parameters
----------
genrep: bool
Whether or not to generate MC replicas. (Only used for checks)
data: validphys.core.DataGroupSpec
containing the datasets to be included in the fit. (Only used
for checks)
experiments_data: list[validphys.core.DataGroupSpec]
similar to `data` but now passed as argument to `ModelTrainer`
replicas_nnseed_fitting_data_dict: list[tuple]
list with element for each replica (typically just one) to be
fitted. Each element
is a tuple containing the replica number, nnseed and
``fitted_data_dict`` containing all of the data, metadata
for each group of datasets which is to be fitted.
posdatasets_fitting_pos_dict: list[dict]
list of dictionaries containing all data and metadata for each
positivity dataset
integdatasets_fitting_integ_dict: list[dict]
list of dictionaries containing all data and metadata for each
integrability dataset
theoryid: validphys.core.TheoryIDSpec
Theory which is used to generate theory predictions from model
during fit. Object also contains some metadata on the theory
settings.
fiatlux: dict
dictionary containing the params needed from LuxQED
basis: list[dict]
preprocessing information for each flavour to be fitted.
fitbasis: str
Valid basis which the fit is to be ran in. Available bases can
be found in :py:mod:`validphys.pdfbases`.
sum_rules: str
Whether to impose sum rules in fit. By default set to True="ALL"
parameters: dict
Mapping containing parameters which define the network
architecture/fitting methodology.
replica_path: pathlib.Path
path to the output of this run
output_path: str
name of the fit
save: None, str
model file where weights will be saved, used in conjunction with
``load``.
load: None, str
model file from which to load weights from.
hyperscanner: dict
dictionary containing the details of the hyperscanner
hyperopt: int
if given, number of hyperopt iterations to run
kfold_parameters: None, dict
dictionary with kfold settings used in hyperopt.
tensorboard: None, dict
mapping containing tensorboard settings if it is to be used. By
default it is None and tensorboard is not enabled.
debug: bool
activate some debug options
maxcores: int
maximum number of (logical) cores that the backend should be aware of
double_precision: bool
whether to use double precision
parallel_models: bool
whether to run models in parallel
"""
from n3fit.backends import set_initial_state
# If debug is active, the initial state will be fixed so that the run is reproducible
set_initial_state(debug=debug, max_cores=maxcores, double_precision=double_precision)
from n3fit.stopwatch import StopWatch
stopwatch = StopWatch()
# All potentially backend dependent imports should come inside the fit function
# so they can eventually be set from the runcard
from n3fit.io.writer import WriterWrapper
from n3fit.model_trainer import ModelTrainer
# Note that this can be run in sequence or in parallel
# To do both cases in the same loop, we uniformize the replica information as:
# - sequential: a list over replicas, each entry containing tuples of length 1
# - parallel: a list of length 1, containing tuples over replicas
#
# Add inner tuples
replicas_info = [
((replica,), (experiment,), (nnseed,))
for replica, experiment, nnseed in replicas_nnseed_fitting_data_dict
]
n_models = len(replicas_info)
if parallel_models:
# Move replicas from outer list to inner tuples
replicas, experiments, nnseeds = [], [], []
for replica, experiment, nnseed in replicas_info:
replicas.extend(replica)
experiments.extend(experiment)
nnseeds.extend(nnseed)
replicas_info = [(tuple(replicas), tuple(experiments), tuple(nnseeds))]
log.info(
"Starting parallel fits from replica %d to %d", replicas[0], replicas[0] + n_models - 1
)
else:
log.info(
"Starting sequential fits from replica %d to %d",
replicas[0],
replicas[0] + n_models - 1,
)
for replica_idxs, exp_info, nnseeds in replicas_info:
log.info("Starting replica fit " + str(replica_idxs))
# Generate a ModelTrainer object
# this object holds all necessary information to train a PDF (up to the NN definition)
the_model_trainer = ModelTrainer(
experiments_data,
exp_info,
posdatasets_fitting_pos_dict,
integdatasets_fitting_integ_dict,
basis,
fitbasis,
nnseeds,
positivity_bound,
debug=debug,
kfold_parameters=kfold_parameters,
max_cores=maxcores,
model_file=load,
sum_rules=sum_rules,
theoryid=theoryid,
lux_params=fiatlux,
replicas=replica_idxs,
)
# This is just to give a descriptive name to the fit function
pdf_gen_and_train_function = the_model_trainer.hyperparametrizable
# Read up the parameters of the NN from the runcard
stopwatch.register_times("replica_set")
########################################################################
# ### Hyperopt #
# If hyperopt is active the parameters of NN will be substituted by the#
# hyoperoptimizable variables. #
# Hyperopt will run for --hyperopt number of iterations before leaving #
# this block #
########################################################################
if hyperopt:
from n3fit.hyper_optimization.hyper_scan import hyper_scan_wrapper
# Note that hyperopt will not run in parallel or with more than one model _for now_
replica_path_set = replica_path / f"replica_{replica_idxs[0]}"
true_best = hyper_scan_wrapper(
replica_path_set, the_model_trainer, hyperscanner, max_evals=hyperopt
)
print("##################")
print("Best model found: ")
for k, i in true_best.items():
print(f" {k} : {i} ")
# In general after we do the hyperoptimization we do not care about the fit
# so just let this die here
break
####################################################################### end of hyperopt
# Ensure hyperopt is off
the_model_trainer.set_hyperopt(False)
# Enable the tensorboard callback
if tensorboard is not None:
profiling = tensorboard.get("profiling", False)
weight_freq = tensorboard.get("weight_freq", 0)
if parallel_models and n_models != 1:
# If using tensorboard when running in parallel
# dump the debugging data to the nnfit folder
replica_path_set = replica_path
else:
replica_path_set = replica_path / f"replica_{replica_idxs[0]}"
log_path = replica_path_set / "tboard"
the_model_trainer.enable_tensorboard(log_path, weight_freq, profiling)
#############################################################################
# ### Fit #
# This function performs the actual fit, it reads all the parameters in the #
# "parameters" dictionary, uses them to generate the NN and trains the net #
#############################################################################
result = pdf_gen_and_train_function(parameters)
stopwatch.register_ref("replica_fitted", "replica_set")
stopping_object = result["stopping_object"]
log.info("Stopped at epoch=%d", stopping_object.stop_epoch)
final_time = stopwatch.stop()
all_chi2s = the_model_trainer.evaluate(stopping_object)
pdf_models = result["pdf_model"].split_replicas()
q0 = theoryid.get_description().get("Q0")
pdf_instances = [N3PDF(pdf_model, fit_basis=basis, Q=q0) for pdf_model in pdf_models]
writer_wrapper = WriterWrapper(
replica_idxs, pdf_instances, stopping_object, all_chi2s, theoryid, final_time
)
writer_wrapper.write_data(replica_path, output_path.name, save)
if tensorboard is not None:
log.info("Tensorboard logging information is stored at %s", log_path)