From ae59c4084eb6930677be67acad2a28922b3045bf Mon Sep 17 00:00:00 2001 From: Henry Wallace Date: Mon, 20 Jan 2025 16:57:10 +0000 Subject: [PATCH] refactor --- src/MaCh3PythonUtils/__init__.py | 4 +- .../apps/__main_mach3_python_utils__.py | 4 +- .../config_reader/config_reader.py | 175 ++++++----- .../file_handling/chain_handler.py | 187 +++++++----- .../fitters/adaption_handler_gpu.py | 30 +- .../fitters/multi_mcmc_gpu.py | 177 +++++++---- .../machine_learning/file_ml_interface.py | 280 ++++++++++-------- .../machine_learning/ml_factory.py | 113 ++++--- .../scikit/scikit_interface.py | 20 +- .../tensorflow/tf_autotune_interface.py | 147 +++++---- .../tensorflow/tf_interface.py | 71 +++-- .../tensorflow/tf_manual_interface.py | 49 +-- .../tensorflow/tf_normalizing_flow_model.py | 58 ++-- .../tensorflow/tf_residual_model.py | 17 +- .../tensorflow/tf_sequential_model.py | 17 +- 15 files changed, 809 insertions(+), 540 deletions(-) diff --git a/src/MaCh3PythonUtils/__init__.py b/src/MaCh3PythonUtils/__init__.py index 9cb47cf..9e7b134 100644 --- a/src/MaCh3PythonUtils/__init__.py +++ b/src/MaCh3PythonUtils/__init__.py @@ -1,2 +1,2 @@ -__version__="2.0.1" -__author__="Henry Wallace" +__version__ = "2.0.1" +__author__ = "Henry Wallace" diff --git a/src/MaCh3PythonUtils/apps/__main_mach3_python_utils__.py b/src/MaCh3PythonUtils/apps/__main_mach3_python_utils__.py index 4a65422..6dd0144 100644 --- a/src/MaCh3PythonUtils/apps/__main_mach3_python_utils__.py +++ b/src/MaCh3PythonUtils/apps/__main_mach3_python_utils__.py @@ -1,7 +1,9 @@ import argparse -def main()->None: + +def main() -> None: from MaCh3PythonUtils.config_reader.config_reader import ConfigReader + parser = argparse.ArgumentParser(usage="python make_plots -c .yaml") parser.add_argument("-c", "--config", help="yaml config file", required=True) diff --git a/src/MaCh3PythonUtils/config_reader/config_reader.py b/src/MaCh3PythonUtils/config_reader/config_reader.py index c4bb76f..8371728 100644 --- a/src/MaCh3PythonUtils/config_reader/config_reader.py +++ b/src/MaCh3PythonUtils/config_reader/config_reader.py @@ -5,15 +5,16 @@ from MaCh3PythonUtils.fitters.multi_mcmc_gpu import MCMCMultGPU from deepmerge import always_merger + class ConfigReader: # Strictly unecessary but nice conceptually _file_handler = None - _interface = None + _interface = None _plot_interface = None - + __default_settings = { # Settings for file and I/O - "FileSettings" : { + "FileSettings": { # Name of input file "FileName": "", # Skip loading a file? Useful for ML @@ -25,134 +26,154 @@ class ConfigReader: # Run an LLH Scan? "RunLLHScan": False, # Run MCMC? - "RunMCMC": False + "RunMCMC": False, }, - - "ParameterSettings":{ - "CircularParameters" : [], + "ParameterSettings": { + "CircularParameters": [], # List of parameter names - "ParameterNames":[], + "ParameterNames": [], # List of cuts - "ParameterCuts":[], + "ParameterCuts": [], # Name of label branch, used in ML "LabelName": "", # Any parameters we want to ignore - "IgnoredParameters":[] + "IgnoredParameters": [], }, - # Specific Settings for ML Applications "MLSettings": { # Name of plots "PlotOutputName": "ml_output.pdf", # Fitter package either SciKit or TensorFlow "FitterPackage": "", - # Fitter Model + # Fitter Model "FitterName": "", # Keyword arguments for fitter - "FitterKwargs" : {}, - #Use an external model that's already been trained? + "FitterKwargs": {}, + # Use an external model that's already been trained? "AddFromExternalModel": False, # Proportion of input data set used for testing (range of 0-1 ) "TestSize": 0.0, # Name to save ML model in - "MLOutputFile": "mlmodel.pkl" + "MLOutputFile": "mlmodel.pkl", }, - # Settings for LLH Scan - "LikelihoodScanSettings": { - "NDivisions": 100 - }, - - + "LikelihoodScanSettings": {"NDivisions": 100}, # Settings for MCMC "MCMCSettings": { "NSteps": 100000, "NChains": 1, "UpdateStep": 100, - "MaxUpdateSteps": 500000 - } + "MaxUpdateSteps": 500000, + }, } - - - + def __init__(self, config: str): """Constructor :param config: Name of yaml config :type config: str - """ - with open(config, 'r') as c: - yaml_config = yaml.safe_load(c) + """ + with open(config, "r") as c: + yaml_config = yaml.safe_load(c) # Update default settings - self.__chain_settings = always_merger.merge(self.__default_settings, yaml_config) - - def make_file_handler(self)->None: - """Sets up file handler object - """ - # Process MCMC chain - self._file_handler = ChainHandler(self.__chain_settings["FileSettings"]["FileName"], - self.__chain_settings["FileSettings"]["ChainName"], - self.__chain_settings["FileSettings"]["Verbose"]) - - self._file_handler.ignore_plots(self.__chain_settings["ParameterSettings"]["IgnoredParameters"]) - self._file_handler.add_additional_plots(self.__chain_settings["ParameterSettings"]["ParameterNames"]) - - self._file_handler.add_additional_plots(self.__chain_settings["ParameterSettings"]["LabelName"], True) - - self._file_handler.add_new_cuts(self.__chain_settings["ParameterSettings"]["ParameterCuts"]) + self.__chain_settings = always_merger.merge( + self.__default_settings, yaml_config + ) + + def make_file_handler(self) -> None: + """Sets up file handler object""" + # Process MCMC chain + self._file_handler = ChainHandler( + self.__chain_settings["FileSettings"]["FileName"], + self.__chain_settings["FileSettings"]["ChainName"], + self.__chain_settings["FileSettings"]["Verbose"], + ) + + self._file_handler.ignore_plots( + self.__chain_settings["ParameterSettings"]["IgnoredParameters"] + ) + self._file_handler.add_additional_plots( + self.__chain_settings["ParameterSettings"]["ParameterNames"] + ) + + self._file_handler.add_additional_plots( + self.__chain_settings["ParameterSettings"]["LabelName"], True + ) + + self._file_handler.add_new_cuts( + self.__chain_settings["ParameterSettings"]["ParameterCuts"] + ) self._file_handler.convert_ttree_to_array() - - def make_ml_interface(self)->None: - """Generates ML interface objects - """ + def make_ml_interface(self) -> None: + """Generates ML interface objects""" if self._file_handler is None: raise Exception("Cannot make interface without opening a file!") - - factory = MLFactory(self._file_handler, self.__chain_settings["ParameterSettings"]["LabelName"], self.__chain_settings["MLSettings"]["PlotOutputName"]) - self._interface = factory.make_interface(self.__chain_settings["MLSettings"]["FitterPackage"], - self.__chain_settings["MLSettings"]["FitterName"], - **self.__chain_settings["MLSettings"]["FitterKwargs"]) - + factory = MLFactory( + self._file_handler, + self.__chain_settings["ParameterSettings"]["LabelName"], + self.__chain_settings["MLSettings"]["PlotOutputName"], + ) + + self._interface = factory.make_interface( + self.__chain_settings["MLSettings"]["FitterPackage"], + self.__chain_settings["MLSettings"]["FitterName"], + **self.__chain_settings["MLSettings"]["FitterKwargs"] + ) + if self.__chain_settings["MLSettings"].get("AddFromExternalModel"): external_model = self.__chain_settings["MLSettings"]["ExternalModel"] external_scaler = self.__chain_settings["MLSettings"]["ExternalScaler"] self._interface.load_model(external_model) self._interface.load_scaler(external_scaler) - + else: - self._interface.set_training_test_set(self.__chain_settings["MLSettings"]["TestSize"]) + self._interface.set_training_test_set( + self.__chain_settings["MLSettings"]["TestSize"] + ) self._interface.train_model() self._interface.test_model() - self._interface.save_model(self.__chain_settings["MLSettings"]["MLOutputFile"]) - self._interface.save_scaler(self.__chain_settings['MLSettings']['MLScalerOutputName']) - + self._interface.save_model( + self.__chain_settings["MLSettings"]["MLOutputFile"] + ) + self._interface.save_scaler( + self.__chain_settings["MLSettings"]["MLScalerOutputName"] + ) + def run_mcmc(self): print("WARNING: MCMC HAS ONLY BEEN TESTED WITH TENSORFLOW INTERFACES!") - mcmc = MCMCMultGPU(self._interface, - self.__chain_settings["MCMCSettings"]["NChains"], - self.__chain_settings["ParameterSettings"]["CircularParameters"], - self.__chain_settings["MCMCSettings"]["UpdateStep"], - self.__chain_settings["MCMCSettings"]["MaxUpdateSteps"]) + mcmc = MCMCMultGPU( + self._interface, + self.__chain_settings["MCMCSettings"]["NChains"], + self.__chain_settings["ParameterSettings"]["CircularParameters"], + self.__chain_settings["MCMCSettings"]["UpdateStep"], + self.__chain_settings["MCMCSettings"]["MaxUpdateSteps"], + ) - mcmc(self.__chain_settings["MCMCSettings"]["NSteps"], - self.__chain_settings["MCMCSettings"]["MCMCOutput"],) + mcmc( + self.__chain_settings["MCMCSettings"]["NSteps"], + self.__chain_settings["MCMCSettings"]["MCMCOutput"], + ) - def __call__(self) -> None: - """Runs over all files from config - """ + """Runs over all files from config""" self.make_file_handler() - + self.make_ml_interface() - if self.__chain_settings["FileSettings"]["RunLLHScan"] and self._interface is not None: - self._interface.run_likelihood_scan(self.__chain_settings["LikelihoodScanSettings"]["NDivisions"]) - - if self.__chain_settings["FileSettings"]["RunMCMC"] and self._interface is not None: - self.run_mcmc() - + if ( + self.__chain_settings["FileSettings"]["RunLLHScan"] + and self._interface is not None + ): + self._interface.run_likelihood_scan( + self.__chain_settings["LikelihoodScanSettings"]["NDivisions"] + ) + if ( + self.__chain_settings["FileSettings"]["RunMCMC"] + and self._interface is not None + ): + self.run_mcmc() diff --git a/src/MaCh3PythonUtils/file_handling/chain_handler.py b/src/MaCh3PythonUtils/file_handling/chain_handler.py index 997b5ca..a2d683a 100644 --- a/src/MaCh3PythonUtils/file_handling/chain_handler.py +++ b/src/MaCh3PythonUtils/file_handling/chain_handler.py @@ -1,6 +1,7 @@ -''' +""" Python tool to load in some generic TTree objects and export to numpy array/pandas dataframe -''' +""" + import uproot as ur import pandas as pd from typing import List, Union, Any @@ -10,6 +11,7 @@ import numpy as np from numpy.typing import NDArray + class ChainHandler: """ Class to load in ROOT files containing a single TTree @@ -19,7 +21,10 @@ class ChainHandler: :param ttree_name: Name of TTree contained in ROOT file :type ttree_name: str, optional """ - def __init__(self, file_name: str, ttree_name: str="posteriors", verbose=False)->None: + + def __init__( + self, file_name: str, ttree_name: str = "posteriors", verbose=False + ) -> None: """_summary_ :param file_name: Input file name @@ -32,64 +37,76 @@ def __init__(self, file_name: str, ttree_name: str="posteriors", verbose=False)- """ print(f"Attempting to open {file_name}") try: - self._posterior_ttree = ur.open(f"{file_name}:{ttree_name}") + self._posterior_ttree = ur.open(f"{file_name}:{ttree_name}") except FileNotFoundError: - raise IOError(f"The file '{file_name}' does not exist or does not contain '{ttree_name}") - - print(f"Succesfully opened {file_name}:{ttree_name}") - warnings.filterwarnings("ignore", category=DeprecationWarning) #Some imports are a little older - warnings.filterwarnings("ignore", category=UserWarning) #Some imports are a little older - warnings.filterwarnings("ignore", category=pd.errors.PerformanceWarning) # Not a fan of being yelled at by pandas + raise IOError( + f"The file '{file_name}' does not exist or does not contain '{ttree_name}" + ) - self._plotting_branches = [] # Filled with branches we want to plot - self._cuts = [] # If we want to apply cuts (can be done later but fastest at load time) + print(f"Succesfully opened {file_name}:{ttree_name}") + warnings.filterwarnings( + "ignore", category=DeprecationWarning + ) # Some imports are a little older + warnings.filterwarnings( + "ignore", category=UserWarning + ) # Some imports are a little older + warnings.filterwarnings( + "ignore", category=pd.errors.PerformanceWarning + ) # Not a fan of being yelled at by pandas + + self._plotting_branches = [] # Filled with branches we want to plot + self._cuts = ( + [] + ) # If we want to apply cuts (can be done later but fastest at load time) + + self._ttree_array = None # For storing the eventual TTree - self._ttree_array = None #For storing the eventual TTree - self._is_file_open = True - + self._verbose = verbose self._ignored_branches = [] self._arviz_tree = None - def close_file(self)->None: - ''' + def close_file(self) -> None: + """ Closes ROOT file, should be called to avoid memory issues! - ''' + """ if not self._is_file_open: self._posterior_ttree.close() self._is_file_open = False @property - def plot_branches(self)->List[str]: - ''' + def plot_branches(self) -> List[str]: + """ Getter for list of useful branches :return: List of branches used in file :rtype: list - ''' + """ return self._plotting_branches @plot_branches.setter - def plot_branches(self, useful_branches: List[str])->None: - ''' + def plot_branches(self, useful_branches: List[str]) -> None: + """ Setter for list of useful branches :param useful_branches: List of branches we want to plot with :type useful_branches: list - ''' + """ if not self._is_file_open: raise Warning("Adding branches after shutting the file has no effect") self._plotting_branches = useful_branches - def add_additional_plots(self, additional_branches: List[str] | str, exact=False)->None: - ''' + def add_additional_plots( + self, additional_branches: List[str] | str, exact=False + ) -> None: + """ To add more branches to the plotting branch list :param additional_branches: List of branches to add to the plotting list :type additional_branches: list - ''' + """ if not self._is_file_open: raise Warning("Adding branches after shutting the file has no effect") @@ -98,75 +115,79 @@ def add_additional_plots(self, additional_branches: List[str] | str, exact=False branch_list = [] for key in self._posterior_ttree.keys(): - - if key in self._ignored_branches: continue - - if any(var in key for var in additional_branches) and not exact: # Not the most efficient but adds variables to our list of variables + + if key in self._ignored_branches: + continue + + if ( + any(var in key for var in additional_branches) and not exact + ): # Not the most efficient but adds variables to our list of variables branch_list.append(key) elif exact and key in additional_branches: branch_list.append(key) self._plotting_branches.extend(branch_list) - - def ignore_plots(self, ignored_branches: List[str]| str)->None: + + def ignore_plots(self, ignored_branches: List[str] | str) -> None: """List of plots to ignore :param ignored_branches: _description_ :type ignored_branches: List[str] | str - """ + """ if isinstance(ignored_branches, str): ignored_branches = list(ignored_branches) - + for branch in ignored_branches: - if branch in self._ignored_branches: continue + if branch in self._ignored_branches: + continue self._ignored_branches.append(branch) - + if branch in self._plotting_branches: self._plotting_branches.remove(branch) - - - def add_new_cuts(self, new_cuts: Union[str, List[str]])->None: - ''' + + def add_new_cuts(self, new_cuts: Union[str, List[str]]) -> None: + """ Specifies list of cuts to apply to the TTree (something like ['step>80000', 'dm23>0']) :param new_cuts: List of/single cut to apply :type new_cuts: list, str - ''' + """ # Hacky but lets us be a little bit polymorphic if not self._is_file_open: raise Warning("Applying cuts after shutting the file has no effect") - if type(new_cuts)==str: + if type(new_cuts) == str: new_cuts = [new_cuts] self._cuts.extend(new_cuts) - def convert_ttree_to_array(self, close_file=True)->None: - ''' + def convert_ttree_to_array(self, close_file=True) -> None: + """ Converts the TTree table to array :param close_file: Do you want to close the ROOT file after calling this method? :type close_file: bool, optional - ''' + """ if not self._is_file_open: raise IOError("Cannot convert TTree to array after input ROOT file is shut") cuts = "" - if len(self._cuts)>0: + if len(self._cuts) > 0: cuts = f"*".join(f"({cut})" for cut in self._cuts) - - + with ThreadPoolExecutor() as executor: # Make sure we have loads of memory available! # Ensures we don't run into funny behaviour when uncompressing - total_memory_needed = self._posterior_ttree.uncompressed_bytes #in bytes + total_memory_needed = self._posterior_ttree.uncompressed_bytes # in bytes if self._verbose: - print(f"Using {executor._max_workers} threads and requiring {np.round(self._posterior_ttree.uncompressed_bytes*1e-9,3)} Gb memory") + print( + f"Using {executor._max_workers} threads and requiring {np.round(self._posterior_ttree.uncompressed_bytes*1e-9,3)} Gb memory" + ) print("Using the following branches: ") for i in self._plotting_branches: print(f" -> {i}") - + # To make sure we don't run into any unpleasantness # total_available_memory = int(psutil.virtual_memory().available) # if total_memory_needed < total_available_memory: # For some reason I can't just do neededNone: # We're going to surpress some pandas warnings here since ROOT isn't totally efficient when converting to Pandas (but it'll do!) # warnings.filterwarnings("ignore", category=pd.errors.PerformanceWarning) # Now we generate an array object - if len(self._plotting_branches)==0: - self._ttree_array = self._posterior_ttree.arrays(self._posterior_ttree.keys(), cut=cuts, library='pd', decompression_executor=executor, interpretation_executor=executor) # Load in ROOT TTree + if len(self._plotting_branches) == 0: + self._ttree_array = self._posterior_ttree.arrays( + self._posterior_ttree.keys(), + cut=cuts, + library="pd", + decompression_executor=executor, + interpretation_executor=executor, + ) # Load in ROOT TTree else: - self._ttree_array = self._posterior_ttree.arrays(self._plotting_branches, cut=cuts, library='pd', array_cache=f"{total_memory_needed} b", decompression_executor=executor, interpretation_executor=executor) # Load in ROOT TTree + self._ttree_array = self._posterior_ttree.arrays( + self._plotting_branches, + cut=cuts, + library="pd", + array_cache=f"{total_memory_needed} b", + decompression_executor=executor, + interpretation_executor=executor, + ) # Load in ROOT TTree if self._verbose: - print(f"Converted TTree to pandas dataframe with {len(self._ttree_array)} elements") + print( + f"Converted TTree to pandas dataframe with {len(self._ttree_array)} elements" + ) if close_file: self.close_file() # Just to really make sure we have no memory overuse we'll call the Garbage collector - gc.collect() - + gc.collect() @property - def ttree_array(self)->pd.DataFrame: - ''' + def ttree_array(self) -> pd.DataFrame: + """ Getter for the converted TTree array :return: Table containing TTree in non-ROOT format :rtype: Union[np.array, pd.DataFrame, ak.Array] - ''' + """ return self._ttree_array - + @ttree_array.setter - def ttree_array(self, new_array: Any=None)->None: - ''' + def ttree_array(self, new_array: Any = None) -> None: + """ Setter for TTree array object ::: NOTE THIS WILL JUST RAISE AN ERROR :param new_array: Object to set our ttree_array_to :type new_array: Any - ''' + """ # Implemented in case someone tries to do something daft! - raise NotImplementedError("Cannot set converted TTree array to new type") - + raise NotImplementedError("Cannot set converted TTree array to new type") + @property - def ndim(self)->int: - if self._ttree_array is None: + def ndim(self) -> int: + if self._ttree_array is None: return 0 return self._ttree_array.shape[1] - + @property - def lower_bounds(self)->NDArray: + def lower_bounds(self) -> NDArray: # Lower bounds for all params - if self._ttree_array is None: + if self._ttree_array is None: return np.empty(self.ndim) return self._ttree_array.min(axis=0).to_numpy() @property - def upper_bounds(self)->NDArray: + def upper_bounds(self) -> NDArray: # Upper bounds for all params - if self._ttree_array is None: + if self._ttree_array is None: return np.empty(self.ndim) - - return self._ttree_array.max(axis=0).to_numpy() \ No newline at end of file + + return self._ttree_array.max(axis=0).to_numpy() diff --git a/src/MaCh3PythonUtils/fitters/adaption_handler_gpu.py b/src/MaCh3PythonUtils/fitters/adaption_handler_gpu.py index 2526676..f578f2c 100644 --- a/src/MaCh3PythonUtils/fitters/adaption_handler_gpu.py +++ b/src/MaCh3PythonUtils/fitters/adaption_handler_gpu.py @@ -1,14 +1,17 @@ import tensorflow as tf + class CovarianceUpdaterGPU: def __init__(self, n_dimensions, update_step=1000, max_update: int = 20000): self.n = n_dimensions self._max_update = max_update # Initialize mean and covariance as TensorFlow tensors self.mean = tf.Variable(tf.zeros(n_dimensions, dtype=tf.float32)) - self.covariance = tf.Variable(tf.zeros((n_dimensions, n_dimensions), dtype=tf.float32)) + self.covariance = tf.Variable( + tf.zeros((n_dimensions, n_dimensions), dtype=tf.float32) + ) self.frozen_covariance = tf.eye(n_dimensions, dtype=tf.float32) - + self.count = 0 self.update_step = update_step @@ -19,8 +22,13 @@ def update_covariance(self, new_data): delta = new_data - self.mean self.mean.assign_add(delta / count) delta2 = new_data - self.mean - self.covariance.assign_add(tf.linalg.matmul(tf.expand_dims(delta, axis=-1), - tf.expand_dims(delta2, axis=0)) * (count - 1) / count) + self.covariance.assign_add( + tf.linalg.matmul( + tf.expand_dims(delta, axis=-1), tf.expand_dims(delta2, axis=0) + ) + * (count - 1) + / count + ) return self.mean, self.covariance @tf.function @@ -31,11 +39,11 @@ def update(self, new_data): new_data (tf.Tensor): A 1D tensor representing the new data point (shape: (n_dimensions,)) """ self.count += 1 - + # Arbitary stopping point! - if self.count>self._max_update: + if self.count > self._max_update: return - + # Update mean and covariance using the class method self.mean, self.covariance = self.update_covariance(new_data) @@ -50,10 +58,14 @@ def get_mean(self): @tf.function def get_covariance(self): """Return the current covariance matrix.""" - return self.covariance / (self.count - 1) if self.count > 1 else tf.eye((self.n), dtype=tf.float32) + return ( + self.covariance / (self.count - 1) + if self.count > 1 + else tf.eye((self.n), dtype=tf.float32) + ) def get_frozen_covariance(self): - return (2.4 ** 2) / float(self.n) * self.frozen_covariance + tf.eye(self.n) * 1e-6 + return (2.4**2) / float(self.n) * self.frozen_covariance + tf.eye(self.n) * 1e-6 @tf.function def sample(self, n_samples: int): diff --git a/src/MaCh3PythonUtils/fitters/multi_mcmc_gpu.py b/src/MaCh3PythonUtils/fitters/multi_mcmc_gpu.py index 09ebe55..09c7f1d 100644 --- a/src/MaCh3PythonUtils/fitters/multi_mcmc_gpu.py +++ b/src/MaCh3PythonUtils/fitters/multi_mcmc_gpu.py @@ -13,34 +13,59 @@ from MaCh3PythonUtils.machine_learning.tensorflow.tf_interface import TfInterface from MaCh3PythonUtils.fitters.adaption_handler_gpu import CovarianceUpdaterGPU + class MCMCMultGPU: - def __init__(self, interface: TfInterface, n_chains: int = 1024, circular_params: List[str] = [], update_step: int = 10, max_update: int = 200000): + def __init__( + self, + interface: TfInterface, + n_chains: int = 1024, + circular_params: List[str] = [], + update_step: int = 10, + max_update: int = 200000, + ): print("MCMC let's go!") - self._interface = interface + self._interface = interface self._n_dim = interface.chain.ndim - 1 self._n_chains = n_chains # Initial states for all chains initial_state = tf.convert_to_tensor(np.zeros(self._n_dim), dtype=tf.float32) - self._chain_states = tf.Variable(tf.tile(tf.expand_dims(initial_state, axis=0), [n_chains, 1]), dtype=tf.float32) + self._chain_states = tf.Variable( + tf.tile(tf.expand_dims(initial_state, axis=0), [n_chains, 1]), + dtype=tf.float32, + ) # boundaries - self._upper_bounds = tf.convert_to_tensor(self._interface.scale_data(self._interface.chain.upper_bounds[:-1].reshape(1,-1)), dtype=tf.float32) - self._lower_bounds = tf.convert_to_tensor(self._interface.scale_data(self._interface.chain.lower_bounds[:-1].reshape(1,-1)), dtype=tf.float32) - + self._upper_bounds = tf.convert_to_tensor( + self._interface.scale_data( + self._interface.chain.upper_bounds[:-1].reshape(1, -1) + ), + dtype=tf.float32, + ) + self._lower_bounds = tf.convert_to_tensor( + self._interface.scale_data( + self._interface.chain.lower_bounds[:-1].reshape(1, -1) + ), + dtype=tf.float32, + ) + max_val = self._interface.chain.upper_bounds[-1] - + self._max_val_tf = tf.constant(max_val) self._circular_indices = self._get_circular_indices(circular_params) print(self._circular_indices) # CovarianceUpdater will be updated based on the first chain - self._matrix_handler = CovarianceUpdaterGPU(self._n_dim, update_step, max_update) + self._matrix_handler = CovarianceUpdaterGPU( + self._n_dim, update_step, max_update + ) # Calculate likelihoods for all chains (with GPU processing) - self._current_loglikelihoods = tf.Variable(self._calc_likelihood(self._chain_states)) + self._current_loglikelihoods = tf.Variable( + self._calc_likelihood(self._chain_states) + ) self._current_step = 0 # Automate batch size determination based on memory availability for this process @@ -50,18 +75,20 @@ def __init__(self, interface: TfInterface, n_chains: int = 1024, circular_params self._queue = tf.queue.FIFOQueue( capacity=self._batch_size_steps, dtypes=[tf.float32], - shapes=[(self._n_chains, self._n_dim)] + shapes=[(self._n_chains, self._n_dim)], ) def _get_circular_indices(self, circular_params: List[str]): """Map circular params to indices in self._interface.chain.plot_branches.""" - return [self._interface.chain.plot_branches.index(param) for param in circular_params] - + return [ + self._interface.chain.plot_branches.index(param) + for param in circular_params + ] def _estimate_batch_size(self): """Estimate batch size based on memory available to this process.""" step_size_in_bytes = self._n_chains * self._n_dim * tf.float32.size - + # Get memory info for the current process process = psutil.Process() available_memory = process.memory_info().rss # Get memory in use by the process @@ -69,7 +96,9 @@ def _estimate_batch_size(self): usable_memory = available_memory * memory_fraction estimated_batch_size = int(usable_memory // step_size_in_bytes) - print(f"Automatically determined batch size for the process: {estimated_batch_size}") + print( + f"Automatically determined batch size for the process: {estimated_batch_size}" + ) return estimated_batch_size @@ -84,34 +113,43 @@ def _calc_likelihood(self, states: tf.Tensor): @tf.function def propose_step_gpu(self): # Propose new states for all chains - proposed_states = self._matrix_handler.sample(self._n_chains) + self._chain_states - + proposed_states = ( + self._matrix_handler.sample(self._n_chains) + self._chain_states + ) + def apply_circular_bounds(idx): # Extract specific bounds for the circular parameter lower_bound = self._lower_bounds[0, idx] upper_bound = self._upper_bounds[0, idx] - adjusted_values = lower_bound + tf.math.mod(proposed_states[:, idx] - upper_bound, upper_bound - lower_bound) + adjusted_values = lower_bound + tf.math.mod( + proposed_states[:, idx] - upper_bound, upper_bound - lower_bound + ) return tf.tensor_scatter_nd_update( proposed_states, indices=[[chain_idx, idx] for chain_idx in range(self._n_chains)], - updates=adjusted_values + updates=adjusted_values, ) # Apply circular bounds to indices marked as circular for idx in self._circular_indices: proposed_states = apply_circular_bounds(idx) - # Apply boundary conditions - proposed_states = tf.where(proposed_states < self._lower_bounds, self._chain_states, proposed_states) - proposed_states = tf.where(proposed_states > self._upper_bounds, self._chain_states, proposed_states) + proposed_states = tf.where( + proposed_states < self._lower_bounds, self._chain_states, proposed_states + ) + proposed_states = tf.where( + proposed_states > self._upper_bounds, self._chain_states, proposed_states + ) # Calculate log-likelihoods for proposed states proposed_loglikelihoods = self._calc_likelihood(proposed_states) # Metropolis-Hastings acceptance step - log_diff = proposed_loglikelihoods - self._current_loglikelihoods # Because network scales itself - + log_diff = ( + proposed_loglikelihoods - self._current_loglikelihoods + ) # Because network scales itself + # Need to undo the scaling done in the NN log_diff = self._max_val_tf * log_diff acc_probs = tf.minimum(1.0, tf.exp(tf.clip_by_value(log_diff, -100, 0))) @@ -124,7 +162,9 @@ def apply_circular_bounds(idx): # Use `tf.where` to conditionally update the chain states based on acceptance self._chain_states.assign(tf.where(accept, proposed_states, self._chain_states)) - self._current_loglikelihoods.assign(tf.where(accept, proposed_loglikelihoods, self._current_loglikelihoods)) + self._current_loglikelihoods.assign( + tf.where(accept, proposed_loglikelihoods, self._current_loglikelihoods) + ) # Update covariance based on first chain self._matrix_handler.update(self._chain_states[0]) @@ -146,7 +186,7 @@ def propose_step(self): def _flush_async(self, final_flush=False): """ Dequeue cached steps in a single batch and write to the HDF5 file. - + If `final_flush` is True, flush the remaining items even if they don't fill a full batch. """ @@ -167,48 +207,53 @@ def _flush_async(self, final_flush=False): steps_to_write = self._queue.dequeue_many(self._batch_size_steps) end_idx = self._current_step - self._dataset[end_idx-len(steps_to_write):end_idx, :] = steps_to_write + self._dataset[end_idx - len(steps_to_write) : end_idx, :] = steps_to_write - - def save_mcmc_chain_to_pdf(self, filename: str, output_pdf: str, thinning: int=10, burnin: int=10000): + def save_mcmc_chain_to_pdf( + self, filename: str, output_pdf: str, thinning: int = 10, burnin: int = 10000 + ): # Open the HDF5 file and read the chain - with h5py.File(filename, 'r') as f: - chain = f['chain'][burnin::thinning] + with h5py.File(filename, "r") as f: + chain = f["chain"][burnin::thinning] # Need it to reflect the actual parameters in our fit so let's combine everything! - rescaled_chain = [self._interface.invert_scaling(chain[1000:,i]) for i in range(self._n_chains)] + rescaled_chain = [ + self._interface.invert_scaling(chain[1000:, i]) + for i in range(self._n_chains) + ] combined_rescaled_chain = np.concatenate(rescaled_chain, axis=0) - + _, n_params = chain.shape[1:] - + # Create a PdfPages object to save plots print("Plotting traces") with PdfPages(output_pdf) as pdf: - + # Rescale the chain - + for i in tqdm(range(n_params)): fig, ax = plt.subplots(figsize=(10, 6)) # Plot the chain for the i-th parameter # unscaled_data = self._interface.invert_scaling(chain[:, 0, i]) # for n, r in enumerate(rescaled_chain): - ax.plot(rescaled_chain[0][:, i], lw=0.5, label=f'Chain 0') + ax.plot(rescaled_chain[0][:, i], lw=0.5, label=f"Chain 0") ax.set_ylabel(self._interface.chain.plot_branches[i]) - ax.set_title(f"Parameter {self._interface.chain.plot_branches[i]} MCMC Chain") - ax.set_xlabel('Step') + ax.set_title( + f"Parameter {self._interface.chain.plot_branches[i]} MCMC Chain" + ) + ax.set_xlabel("Step") # Save the current figure to the PDF pdf.savefig(fig) plt.close(fig) # Close the figure to save memory - # Create a PdfPages object to save plots print("Plotting posteriors") with PdfPages(f"posterior_{output_pdf}") as pdf: - + # Rescale the chain - + for i in tqdm(range(n_params)): fig, ax = plt.subplots(figsize=(10, 6)) @@ -217,12 +262,28 @@ def save_mcmc_chain_to_pdf(self, filename: str, output_pdf: str, thinning: int=1 l = self._interface.chain.lower_bounds[i] u = self._interface.chain.upper_bounds[i] bins = np.linspace(l, u, 100) - - ax.hist(rescaled_chain[0][:, i], color='b', label="ML Pred", alpha=0.3, bins=bins, density=True) - ax.hist(self._interface.test_data.iloc[10000:,i].to_numpy(), color='r', label="Real Result", alpha=0.3, bins=bins, density=True) - + + ax.hist( + rescaled_chain[0][:, i], + color="b", + label="ML Pred", + alpha=0.3, + bins=bins, + density=True, + ) + ax.hist( + self._interface.test_data.iloc[10000:, i].to_numpy(), + color="r", + label="Real Result", + alpha=0.3, + bins=bins, + density=True, + ) + ax.set_xlabel(self._interface.chain.plot_branches[i]) - ax.set_title(f"Parameter {self._interface.chain.plot_branches[i]} MCMC Chain") + ax.set_title( + f"Parameter {self._interface.chain.plot_branches[i]} MCMC Chain" + ) ax.legend() # Save the current figure to the PDF @@ -237,11 +298,15 @@ def save_mcmc_chain_to_pdf(self, filename: str, output_pdf: str, thinning: int=1 # Plot the chain for the i-th parameter # unscaled_data = self._interface.invert_scaling(chain[:, 0, i]) # for n, r in enumerate(rescaled_chain): - ac = sm.tsa.acf(rescaled_chain[0][:, i], nlags=len(rescaled_chain[0][:, 1])) - ax.plot(ac, lw=0.5, label=f'Chain 0') + ac = sm.tsa.acf( + rescaled_chain[0][:, i], nlags=len(rescaled_chain[0][:, 1]) + ) + ax.plot(ac, lw=0.5, label=f"Chain 0") ax.set_ylabel(self._interface.chain.plot_branches[i]) - ax.set_title(f"Parameter {self._interface.chain.plot_branches[i]} MCMC Chain") - ax.set_xlabel('Autocorrelation') + ax.set_title( + f"Parameter {self._interface.chain.plot_branches[i]} MCMC Chain" + ) + ax.set_xlabel("Autocorrelation") # Save the current figure to the PDF pdf.savefig(fig) @@ -253,12 +318,16 @@ def __call__(self, n_steps, output_file_name: str): print(f"Running MCMC for {n_steps} steps with {self._n_chains} chains") # Open the HDF5 file in append mode - with h5py.File(output_file_name, 'w') as f: + with h5py.File(output_file_name, "w") as f: # Create or resize the dataset - if 'chain' in f: - del f['chain'] # Delete if it already exists to avoid appending duplicate data + if "chain" in f: + del f[ + "chain" + ] # Delete if it already exists to avoid appending duplicate data - self._dataset = f.create_dataset('chain', (n_steps, self._n_chains, self._n_dim), chunks=True) + self._dataset = f.create_dataset( + "chain", (n_steps, self._n_chains, self._n_dim), chunks=True + ) for _ in tqdm(range(n_steps)): self.propose_step() diff --git a/src/MaCh3PythonUtils/machine_learning/file_ml_interface.py b/src/MaCh3PythonUtils/machine_learning/file_ml_interface.py index 793ab3b..7860d5c 100644 --- a/src/MaCh3PythonUtils/machine_learning/file_ml_interface.py +++ b/src/MaCh3PythonUtils/machine_learning/file_ml_interface.py @@ -20,18 +20,25 @@ import matplotlib.pyplot as plt from matplotlib.backends.backend_pdf import PdfPages + class FileMLInterface(ABC): - white_viridis = LinearSegmentedColormap.from_list('white_viridis', [ - (0, '#ffffff'), - (1e-20, '#440053'), - (0.2, '#404388'), - (0.3, '#2a788e'), - (0.4, '#21a784'), - (0.7, '#78d151'), - (1, '#fde624'), - ], N=256) - - def __init__(self, chain: ChainHandler, prediction_variable: str, fit_name: str) -> None: + white_viridis = LinearSegmentedColormap.from_list( + "white_viridis", + [ + (0, "#ffffff"), + (1e-20, "#440053"), + (0.2, "#404388"), + (0.3, "#2a788e"), + (0.4, "#21a784"), + (0.7, "#78d151"), + (1, "#fde624"), + ], + N=256, + ) + + def __init__( + self, chain: ChainHandler, prediction_variable: str, fit_name: str + ) -> None: """General Interface for all ML models :param chain: ChainHandler instance @@ -39,56 +46,59 @@ def __init__(self, chain: ChainHandler, prediction_variable: str, fit_name: str) :param prediction_variable: "Label" used for prediction :type prediction_variable: str :raises ValueError: Checks to see if label exists in tree - """ + """ self._chain = chain - + self._fit_name = fit_name self._prediction_variable = prediction_variable - + if prediction_variable not in self._chain.ttree_array.columns: raise ValueError(f"Cannot find {prediction_variable} in input tree") - + self._model = None - - self._training_data=None - self._training_labels=None - self._test_data=None - self._test_labels=None + + self._training_data = None + self._training_labels = None + self._test_data = None + self._test_labels = None # Scaling components self._scaler = StandardScaler() # self._pca_matrix = PCA(n_components=0.95) - + self._label_scaler = MinMaxScaler(feature_range=(0, 1)) - - - - def __separate_dataframe(self)->Tuple[pd.DataFrame, pd.DataFrame]: + + def __separate_dataframe(self) -> Tuple[pd.DataFrame, pd.DataFrame]: """Split data frame into feature + label objects :return: features, labels :rtype: Tuple[pd.DataFrame, pd.DataFrame] - """ + """ # Separates dataframe into features + labels features = self._chain.ttree_array.copy() - labels = pd.DataFrame(features.pop(self._prediction_variable) ) - + labels = pd.DataFrame(features.pop(self._prediction_variable)) + return features, labels - + def set_training_test_set(self, test_size: float): """Splits data/labels into training and testing tests :param test_size: Proportion of data used for testing :type test_size: float - """ + """ # Splits in traing + test_spit features, labels = self.__separate_dataframe() - self._training_data, self._test_data, self._training_labels, self._test_labels = train_test_split(features, labels, test_size=test_size) + ( + self._training_data, + self._test_data, + self._training_labels, + self._test_labels, + ) = train_test_split(features, labels, test_size=test_size) # Fit scaling pre-processors. These get applied properly when scale_data is called - _= self._scaler.fit_transform(self._training_data) + _ = self._scaler.fit_transform(self._training_data) self._label_scaler.fit_transform(self._training_labels) - + # self._pca_matrix.fit(scaled_training) def scale_data(self, input_data): @@ -96,7 +106,7 @@ def scale_data(self, input_data): scale_data = self._scaler.transform(input_data) # scale_data = self._pca_matrix.transform(scale_data) return scale_data - + def scale_labels(self, labels): return self._label_scaler.transform(labels) @@ -107,110 +117,105 @@ def invert_scaling(self, input_data): return unscaled_data @property - def model(self)->Any: + def model(self) -> Any: """Model used :return: Returns ML model being used :rtype: Any - """ + """ # Returns model being used - return self._model - + return self._model + @property - def chain(self)->ChainHandler: + def chain(self) -> ChainHandler: return self._chain - + @property - def training_data(self)->pd.DataFrame: + def training_data(self) -> pd.DataFrame: """Gets training data :return: Training data set :rtype: pd.DataFrame - """ + """ if self._training_data is None: - return self._chain.ttree_array.iloc[:,:-1] + return self._chain.ttree_array.iloc[:, :-1] return self._training_data @property - def test_data(self)->pd.DataFrame: + def test_data(self) -> pd.DataFrame: """Gets training data :return: Training data set :rtype: pd.DataFrame - """ + """ if self._test_data is None: - return self._chain.ttree_array.iloc[:,:-1] - + return self._chain.ttree_array.iloc[:, :-1] + return self._test_data - - def add_model(self, ml_model: Any)->None: + def add_model(self, ml_model: Any) -> None: """Add model to data set :param ml_model: Sets model to be ml_model :type ml_model: Any - """ - # Add ML model into your interface + """ + # Add ML model into your interface self._model = ml_model - + @abstractmethod def train_model(self): - """Abstract method, should be overwritten with model training - """ + """Abstract method, should be overwritten with model training""" # Train Model method pass - + @abstractmethod - def model_predict(self, testing_data: pd.DataFrame)->Iterable: + def model_predict(self, testing_data: pd.DataFrame) -> Iterable: """Abstract method, should return model prediction - :param testing_data: Data to test model on + :param testing_data: Data to test model on :type testing_data: pd.DataFrame """ pass - + def save_model(self, output_file: str): """Save model to pickle :param output_file: Pickle file to save to :type output_file: str - """ + """ print(f"Saving to {output_file}") - with open(output_file, 'wb') as f: + with open(output_file, "wb") as f: pickle.dump(self._model, f) def save_scaler(self, output_file: str): - pickle.dump(self._scaler, open(output_file, 'wb')) - + pickle.dump(self._scaler, open(output_file, "wb")) + def load_scaler(self, input_scaler: str): - self._scaler = pickle.load(open(input_scaler, 'rb')) + self._scaler = pickle.load(open(input_scaler, "rb")) - def load_model(self, input_model: str): """Unpickle model :param input_file: Pickled Model :type input_file: str - """ + """ print(f"Attempting to load file from {input_file}") - with open(input_model, 'r') as f: + with open(input_model, "r") as f: self._model = pickle.load(f) - - + def test_model(self): """Test model :raises ValueError: No model set - :raises ValueError: No test data set - """ + :raises ValueError: No test data set + """ if self._model is None: raise ValueError("No Model has been set!") if self._test_data is None or self._test_labels is None: raise ValueError("No test data set") - print("Training Results!") train_prediction = self.model_predict(self._training_data) train_as_numpy = self.scale_labels(self._training_labels).T[0] @@ -221,49 +226,56 @@ def test_model(self): test_prediction = self.model_predict(self._test_data) test_as_numpy = self.scale_labels(self._test_labels).T[0] - + self.evaluate_model(test_prediction, test_as_numpy, outfile=f"{self._fit_name}") print("=====\n\n") - + def model_predict_single_sample(self, sample): - sample_shaped = sample.reshape(1,-1) + sample_shaped = sample.reshape(1, -1) return self.model_predict(sample_shaped)[0] - - def get_maxlikelihood(self)->OptimizeResult: + + def get_maxlikelihood(self) -> OptimizeResult: init_vals = self.training_data.iloc[[1]].to_numpy()[0] - + print("Calculating max LLH") - maximal_likelihood = minimize(self.model_predict_single_sample, init_vals, bounds=zip(self._chain.lower_bounds[:-1], self._chain.upper_bounds[:-1]), method="L-BFGS-B", options={"disp": True}) + maximal_likelihood = minimize( + self.model_predict_single_sample, + init_vals, + bounds=zip(self._chain.lower_bounds[:-1], self._chain.upper_bounds[:-1]), + method="L-BFGS-B", + options={"disp": True}, + ) return maximal_likelihood - def run_likelihood_scan(self, n_divisions: int = 500): # Get nominals print("Running LLH Scan") maximal_likelihood = self.get_maxlikelihood() - maximal_nominal=maximal_likelihood.x - - errors = np.sqrt(np.diag(maximal_likelihood.hess_inv(np.identity(self.chain.ndim-1)))) + maximal_nominal = maximal_likelihood.x - print("Maximal Pars :") - for i in range(self.chain.ndim-1): - print(f"Param : {self.chain.plot_branches[i]} : {maximal_likelihood.x[i]}±{errors[i]}") + errors = np.sqrt( + np.diag(maximal_likelihood.hess_inv(np.identity(self.chain.ndim - 1))) + ) + print("Maximal Pars :") + for i in range(self.chain.ndim - 1): + print( + f"Param : {self.chain.plot_branches[i]} : {maximal_likelihood.x[i]}±{errors[i]}" + ) with PdfPages("llh_scan.pdf") as pdf: - for i in tqdm(range(self.chain.ndim-1), total=self.chain.ndim-1): + for i in tqdm(range(self.chain.ndim - 1), total=self.chain.ndim - 1): # Make copy since we'll be modifying! - + lower_bound = self.chain.lower_bounds[i] upper_bound = self.chain.upper_bounds[i] - + param_range = np.linspace(lower_bound, upper_bound, n_divisions) modified_values = [maximal_nominal.copy() for _ in range(n_divisions)] - - + for j, div in enumerate(param_range): - modified_values[j][i]=div - + modified_values[j][i] = div + prediction = self.model_predict(modified_values) # Save as histogram plt.plot(param_range, prediction) @@ -271,10 +283,10 @@ def run_likelihood_scan(self, n_divisions: int = 500): plt.ylabel("-2*loglikelihood") pdf.savefig() plt.close() - - - - def evaluate_model(self, predicted_values: Iterable, true_values: Iterable, outfile: str=""): + + def evaluate_model( + self, predicted_values: Iterable, true_values: Iterable, outfile: str = "" + ): """Evalulates model :param predicted_values: Label values predicted by model @@ -283,56 +295,76 @@ def evaluate_model(self, predicted_values: Iterable, true_values: Iterable, outf :type true_values: Iterable :param outfile: File to output plots to, defaults to "" :type outfile: str, optional - """ - + """ + print(predicted_values) - print(f"Mean Absolute Error : {metrics.mean_absolute_error(predicted_values,true_values)}") - - + print( + f"Mean Absolute Error : {metrics.mean_absolute_error(predicted_values,true_values)}" + ) + lobf = np.poly1d(np.polyfit(predicted_values, true_values, 1)) - + print(f"Line of best fit : y={lobf.c[0]}x + {lobf.c[1]}") - + fig = plt.figure() - - - ax = fig.add_subplot(1,1,1, projection='scatter_density') - - # Bit hacky put plotting code is... bad so we're going to ignore the error it raises! + + ax = fig.add_subplot(1, 1, 1, projection="scatter_density") + + # Bit hacky put plotting code is... bad so we're going to ignore the error it raises! warnings.filterwarnings("ignore", message="All-NaN slice encountered") - density = ax.scatter_density(predicted_values, true_values, cmap=self.white_viridis) + density = ax.scatter_density( + predicted_values, true_values, cmap=self.white_viridis + ) # warnings.resetwarnings() - + fig.colorbar(density, label="number of points per pixel") - + lims = [ np.min([ax.get_xlim(), ax.get_ylim()]), # min of both axes np.max([ax.get_xlim(), ax.get_ylim()]), # max of both axes ] - ax.plot(lims, lobf(lims), "m", label=f"Best fit: true={lobf.c[0]}pred + {lobf.c[1]}", linestyle="dashed", linewidth=0.3) - - ax.plot(lims, lims, 'r', alpha=0.75, zorder=0, label="true=predicted", linestyle="dashed", linewidth=0.3) - ax.set_aspect('equal') + ax.plot( + lims, + lobf(lims), + "m", + label=f"Best fit: true={lobf.c[0]}pred + {lobf.c[1]}", + linestyle="dashed", + linewidth=0.3, + ) + + ax.plot( + lims, + lims, + "r", + alpha=0.75, + zorder=0, + label="true=predicted", + linestyle="dashed", + linewidth=0.3, + ) + ax.set_aspect("equal") ax.set_xlim(lims) ax.set_ylim(lims) - ax.set_xlabel("Predicted Log likelihood") ax.set_ylabel("True Log Likelihood") - + fig.legend() - if outfile=="": outfile = f"evaluated_model_qq_tf.pdf" - + if outfile == "": + outfile = f"evaluated_model_qq_tf.pdf" + print(f"Saving QQ to {outfile}") - + fig.savefig(outfile) plt.close() - + # Gonna draw a hist - difs = true_values-predicted_values + difs = true_values - predicted_values print(f"mean: {np.mean(difs)}, std dev: {np.std(difs)}") - plt.hist(difs, bins=100, density=True, range=(np.std(difs)*-5, np.std(difs)*5)) + plt.hist( + difs, bins=100, density=True, range=(np.std(difs) * -5, np.std(difs) * 5) + ) plt.xlabel("True - Pred") plt.savefig(f"diffs_5sigma_range_{outfile}") - plt.close() \ No newline at end of file + plt.close() diff --git a/src/MaCh3PythonUtils/machine_learning/ml_factory.py b/src/MaCh3PythonUtils/machine_learning/ml_factory.py index 39a0d61..3fc671b 100644 --- a/src/MaCh3PythonUtils/machine_learning/ml_factory.py +++ b/src/MaCh3PythonUtils/machine_learning/ml_factory.py @@ -3,12 +3,22 @@ """ from MaCh3PythonUtils.machine_learning.scikit.scikit_interface import SciKitInterface -from MaCh3PythonUtils.machine_learning.tensorflow.tf_autotune_interface import TfAutotuneInterface -from MaCh3PythonUtils.machine_learning.tensorflow.tf_sequential_model import TfSequentialModel -from MaCh3PythonUtils.machine_learning.tensorflow.tf_residual_model import TfResidualModel -from MaCh3PythonUtils.machine_learning.tensorflow.tf_normalizing_flow_model import TfNormalizingFlowModel - -from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import TfManualLayeredInterface +from MaCh3PythonUtils.machine_learning.tensorflow.tf_autotune_interface import ( + TfAutotuneInterface, +) +from MaCh3PythonUtils.machine_learning.tensorflow.tf_sequential_model import ( + TfSequentialModel, +) +from MaCh3PythonUtils.machine_learning.tensorflow.tf_residual_model import ( + TfResidualModel, +) +from MaCh3PythonUtils.machine_learning.tensorflow.tf_normalizing_flow_model import ( + TfNormalizingFlowModel, +) + +from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import ( + TfManualLayeredInterface, +) from MaCh3PythonUtils.machine_learning.tensorflow.tf_interface import TfInterface from MaCh3PythonUtils.file_handling.chain_handler import ChainHandler @@ -20,33 +30,34 @@ class MLFactory: # Implement algorithms here __IMPLEMENTED_ALGORITHMS = { - "scikit" : { - "randomforest" : ske.RandomForestRegressor, - "gradientboost" : ske.GradientBoostingRegressor, - "adaboost" : ske.AdaBoostRegressor, - "histboost" : ske.HistGradientBoostingRegressor + "scikit": { + "randomforest": ske.RandomForestRegressor, + "gradientboost": ske.GradientBoostingRegressor, + "adaboost": ske.AdaBoostRegressor, + "histboost": ske.HistGradientBoostingRegressor, }, "tensorflow": { - "sequential" : TfSequentialModel, + "sequential": TfSequentialModel, "residual": TfResidualModel, "normalizing_flow": TfNormalizingFlowModel, - "autotune": TfAutotuneInterface + "autotune": TfAutotuneInterface, }, } - def __init__(self, input_chain: ChainHandler, prediction_variable: str, plot_name: str): + def __init__( + self, input_chain: ChainHandler, prediction_variable: str, plot_name: str + ): """Constructor for ML factory method :param input_chain: ChainHandler instance :type input_chain: ChainHandler :param prediction_variable: Variable we want to predict the value of :type prediction_variable: str - """ + """ # Common chain across all instances of factory self._chain = input_chain self._prediction_variable = prediction_variable self._plot_name = plot_name - def __setup_package_factory(self, package: str, algorithm: str, **kwargs): """Basic method for initialising factory method @@ -59,68 +70,82 @@ def __setup_package_factory(self, package: str, algorithm: str, **kwargs): :raises ValueError: Algorithm not Implemented :return: Model initialised with kwargs :rtype: Any - """ - - package = package.lower() + """ + + package = package.lower() if package not in self.__IMPLEMENTED_ALGORITHMS.keys(): - raise ValueError(f"{package} not included, currently accepted packages are :\n \ - {list(self.__IMPLEMENTED_ALGORITHMS.keys())}") - + raise ValueError( + f"{package} not included, currently accepted packages are :\n \ + {list(self.__IMPLEMENTED_ALGORITHMS.keys())}" + ) + algorithm = algorithm.lower() - + if algorithm not in self.__IMPLEMENTED_ALGORITHMS[package].keys(): - raise ValueError(f"{algorithm} not implemented for {package}, currently accepted algorithms for {package} are:\n \ - {list(self.__IMPLEMENTED_ALGORITHMS[package].keys())}") + raise ValueError( + f"{algorithm} not implemented for {package}, currently accepted algorithms for {package} are:\n \ + {list(self.__IMPLEMENTED_ALGORITHMS[package].keys())}" + ) return self.__IMPLEMENTED_ALGORITHMS[package][algorithm](**kwargs) - def __make_scikit_model(self, algorithm: str, **kwargs)->SciKitInterface: + def __make_scikit_model(self, algorithm: str, **kwargs) -> SciKitInterface: """Generates scikit model instance :param algorithm: Algorithm from scikit :type algorithm: str :return: SciKitInterface wrapper around model :rtype: SciKitInterface - """ + """ # Simple wrapper for scikit packages - interface = SciKitInterface(self._chain, self._prediction_variable, self._plot_name) - interface.add_model(self.__setup_package_factory(package="scikit", algorithm=algorithm, **kwargs)) + interface = SciKitInterface( + self._chain, self._prediction_variable, self._plot_name + ) + interface.add_model( + self.__setup_package_factory( + package="scikit", algorithm=algorithm, **kwargs + ) + ) + + return interface - return interface - - - def __make_tensorflow_layered_model(self, interface: TfManualLayeredInterface, layers: dict)->TfManualLayeredInterface: + def __make_tensorflow_layered_model( + self, interface: TfManualLayeredInterface, layers: dict + ) -> TfManualLayeredInterface: for layer in layers: - layer_id = list(layer.keys())[0] + layer_id = list(layer.keys())[0] interface.add_layer(layer_id, layer[layer_id]) return interface - def __make_tensorflow_model(self, algorithm: str, **kwargs)->TfInterface: - model_func = self.__IMPLEMENTED_ALGORITHMS["tensorflow"].get(algorithm.lower(), None) - + def __make_tensorflow_model(self, algorithm: str, **kwargs) -> TfInterface: + model_func = self.__IMPLEMENTED_ALGORITHMS["tensorflow"].get( + algorithm.lower(), None + ) + if model_func is None: raise Exception(f"Cannot find {algorithm}") - - model: TfInterface = model_func(self._chain, self._prediction_variable, self._plot_name) + + model: TfInterface = model_func( + self._chain, self._prediction_variable, self._plot_name + ) # Ugh - if algorithm=="sequential" or algorithm=="residual": + if algorithm == "sequential" or algorithm == "residual": print("HERE") model = self.__make_tensorflow_layered_model(model, kwargs["Layers"]) model.set_training_settings(kwargs.get("FitSettings")) - model.build_model(**kwargs["BuildSettings"]) - + return model def make_interface(self, interface_type: str, algorithm: str, **kwargs): interface_type = interface_type.lower() - match(interface_type): + match (interface_type): case "scikit": return self.__make_scikit_model(algorithm, **kwargs) case "tensorflow": return self.__make_tensorflow_model(algorithm, **kwargs) case _: - raise Exception(f"{interface_type} not implemented!") \ No newline at end of file + raise Exception(f"{interface_type} not implemented!") diff --git a/src/MaCh3PythonUtils/machine_learning/scikit/scikit_interface.py b/src/MaCh3PythonUtils/machine_learning/scikit/scikit_interface.py index 9d565a9..1194f74 100644 --- a/src/MaCh3PythonUtils/machine_learning/scikit/scikit_interface.py +++ b/src/MaCh3PythonUtils/machine_learning/scikit/scikit_interface.py @@ -6,25 +6,26 @@ - Add staged predict """ -class SciKitInterface(FileMLInterface): + +class SciKitInterface(FileMLInterface): def train_model(self): """Trains model :raises ValueError: Model not initialised :raises ValueError: Data set not initialised - """ + """ print(f"Training Model") scaled_data = self.scale_data(self._training_data) - + if self._model is None: raise ValueError("No Model has been set!") - + if self._training_data is None or self._training_labels is None: raise ValueError("No test data set") - + self._model.fit(scaled_data, self._training_labels) - - def model_predict(self, test_data: DataFrame)->list: + + def model_predict(self, test_data: DataFrame) -> list: """Gets model prediction :param test_data: Data to predict @@ -32,11 +33,10 @@ def model_predict(self, test_data: DataFrame)->list: :raises ValueError: No model set :return: Model prediction for test_data :rtype: list - """ + """ scale_data = self.scale_data(test_data) - + if self._model is None: raise ValueError("No Model has been set!") return self._model.predict(scale_data) - diff --git a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_autotune_interface.py b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_autotune_interface.py index 9e6012e..0314ec4 100644 --- a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_autotune_interface.py +++ b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_autotune_interface.py @@ -14,91 +14,120 @@ class TfAutotuneInterface(TfInterface): _epochs = 1000 _val_split = 0.2 - + def build_model(self, **kwargs): # Set up layers n_layers = kwargs.get("n_layers", [5, 20, 1]) - - + activation_functions = kwargs.get("layer_activation", ["tanh", "relu"]) - + neurons_per_layer = kwargs.get("neurons_per_layer", [16, 2048, 100]) - learning_rate = kwargs.get("learning_rate", [1, 1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6]) - + learning_rate = kwargs.get( + "learning_rate", [1, 1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6] + ) + regularization_rate = kwargs.get("regularization", [0.001, 0.01, 0.1, 1]) - + # Actually make the model def model_builder(hp): model = tfk.Sequential() - - hp_model_layers = hp.Int('n_layers', min_value = n_layers[0], - max_value=n_layers[1], - step=n_layers[2]) - + + hp_model_layers = hp.Int( + "n_layers", + min_value=n_layers[0], + max_value=n_layers[1], + step=n_layers[2], + ) + for i in range(hp_model_layers): # Add layer - hp_layer_units = hp.Int(f'units_{i}', min_value = neurons_per_layer[0], - max_value=neurons_per_layer[1], - step=neurons_per_layer[2]) - hp_layer_activation = hp.Choice(f'activation_{i}', values=activation_functions) - - hp_layer_regularization =hp.Choice(f'regularization_{i}', regularization_rate) - - model.add(tfk.layers.Dense(units = hp_layer_units, - activation = hp_layer_activation, - kernel_regularizer = tfk.regularizers.L2(hp_layer_regularization))) - + hp_layer_units = hp.Int( + f"units_{i}", + min_value=neurons_per_layer[0], + max_value=neurons_per_layer[1], + step=neurons_per_layer[2], + ) + hp_layer_activation = hp.Choice( + f"activation_{i}", values=activation_functions + ) + + hp_layer_regularization = hp.Choice( + f"regularization_{i}", regularization_rate + ) + + model.add( + tfk.layers.Dense( + units=hp_layer_units, + activation=hp_layer_activation, + kernel_regularizer=tfk.regularizers.L2(hp_layer_regularization), + ) + ) + # Add batch normalization model.add(tfk.layers.BatchNormalization()) - + # Add output model.add(tfk.layers.Dense(units=1, activation="linear")) - hp_learning_rate = hp.Choice('learning_rate', learning_rate) - - model.compile(optimizer=tfk.optimizers.Adam(learning_rate=hp_learning_rate), - loss=kwargs.get('loss', 'mse'), - metrics=kwargs.get('metrics', ['mse', 'mae']) + hp_learning_rate = hp.Choice("learning_rate", learning_rate) + + model.compile( + optimizer=tfk.optimizers.Adam(learning_rate=hp_learning_rate), + loss=kwargs.get("loss", "mse"), + metrics=kwargs.get("metrics", ["mse", "mae"]), ) return model - - - self._epochs = kwargs.get('epochs', 1000) - self._val_split = kwargs.get('validation_split', 0.2) - self._batch_size = kwargs.get('batch_size', 2048) - + + self._epochs = kwargs.get("epochs", 1000) + self._val_split = kwargs.get("validation_split", 0.2) + self._batch_size = kwargs.get("batch_size", 2048) + # Hyperband parameters hyperband_iterations = kwargs.get("hyperband_iterations", 100) model_directory = kwargs.get("tuning_dir", "tuning-model") project_name = kwargs.get("project_name", "tuning-project") - - self._model_tuner= kt.Hyperband(model_builder, - objective='val_loss', - max_epochs=self._epochs, - hyperband_iterations=hyperband_iterations, - directory=model_directory, - # distribution_strategy=tf.distribute.MirroredStrategy(), - project_name=project_name, - overwrite=False,) + + self._model_tuner = kt.Hyperband( + model_builder, + objective="val_loss", + max_epochs=self._epochs, + hyperband_iterations=hyperband_iterations, + directory=model_directory, + # distribution_strategy=tf.distribute.MirroredStrategy(), + project_name=project_name, + overwrite=False, + ) def train_model(self): - + scaled_data = self.scale_data(self._training_data) scaled_labels = self.scale_labels(self._training_labels) - - stop_early = tfk.callbacks.EarlyStopping(monitor='val_loss', patience=5) - lr_schedule = tfk.callbacks.ReduceLROnPlateau(monitor="loss", patience=5, factor=0.5, min_lr=1e-6, verbose=1) - - self._model_tuner.search(scaled_data, scaled_labels, - epochs= self._epochs, - validation_split=self._val_split, - batch_size=self._batch_size, - callbacks=[stop_early]) - - best_hps=self._model_tuner.get_best_hyperparameters()[0] + + stop_early = tfk.callbacks.EarlyStopping(monitor="val_loss", patience=5) + lr_schedule = tfk.callbacks.ReduceLROnPlateau( + monitor="loss", patience=5, factor=0.5, min_lr=1e-6, verbose=1 + ) + + self._model_tuner.search( + scaled_data, + scaled_labels, + epochs=self._epochs, + validation_split=self._val_split, + batch_size=self._batch_size, + callbacks=[stop_early], + ) + + best_hps = self._model_tuner.get_best_hyperparameters()[0] print("Finished auto-tuning") - + self._model = self._model_tuner.hypermodel.build(best_hps) - - self._model.fit(scaled_data, scaled_labels, epochs=self._epochs, batch_size=self._batch_size , validation_split=0.2, callbacks=[stop_early, lr_schedule]) + + self._model.fit( + scaled_data, + scaled_labels, + epochs=self._epochs, + batch_size=self._batch_size, + validation_split=0.2, + callbacks=[stop_early, lr_schedule], + ) diff --git a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_interface.py b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_interface.py index 06f4d04..3ff94c5 100644 --- a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_interface.py +++ b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_interface.py @@ -12,12 +12,12 @@ class TfInterface(FileMLInterface): __TF_LAYER_IMPLEMENTATIONS = { "dense": tfk.layers.Dense, "dropout": tfk.layers.Dropout, - "batchnorm": tfk.layers.BatchNormalization + "batchnorm": tfk.layers.BatchNormalization, } - + _layers = [] _training_settings = {} - + def add_layer(self, layer_id: str, layer_args: dict): """Add new layer to TF model @@ -26,7 +26,7 @@ def add_layer(self, layer_id: str, layer_args: dict): :param layer_args: kwargs for layer :type layer_args: dict :raises ValueError: Layer type not implemented in __TF_LAYER_IMPLEMENTATIONS yet - """ + """ if layer_id not in self.__TF_LAYER_IMPLEMENTATIONS.keys(): raise ValueError(f"{layer_id} not implemented yet!") @@ -34,41 +34,50 @@ def add_layer(self, layer_id: str, layer_args: dict): # Hacky, swaps string value of regularliser for proper one layer_args["kernel_regularizer"] = tfk.regularizers.L2(0.2) - self._layers.append(self.__TF_LAYER_IMPLEMENTATIONS[layer_id.lower()](**layer_args)) - + self._layers.append( + self.__TF_LAYER_IMPLEMENTATIONS[layer_id.lower()](**layer_args) + ) + def build_model(self, **kwargs): return None - def set_training_settings(self, kwargs): """Set training settings, needs to be done early for...reasons :param kwargs: training kwargs :type kwargs: kwargs - """ - self._training_settings = kwargs - + """ + self._training_settings = kwargs + def train_model(self): - """train model - """ + """train model""" scaled_data = self.scale_data(self._training_data) - - lr_schedule = tfk.callbacks.ReduceLROnPlateau(monitor="loss", patience=5, factor=0.5, min_lr=1e-6, verbose=1) - self._model.fit(scaled_data, self._training_labels, **self._training_settings, callbacks=[lr_schedule]) - print(f"Using loss function: {self._model.loss}") + lr_schedule = tfk.callbacks.ReduceLROnPlateau( + monitor="loss", patience=5, factor=0.5, min_lr=1e-6, verbose=1 + ) + + self._model.fit( + scaled_data, + self._training_labels, + **self._training_settings, + callbacks=[lr_schedule], + ) + print(f"Using loss function: {self._model.loss}") def save_model(self, output_file: str): """Save model to file - :param output_file: Output file to save model to - :type output_file: str - """ + :param output_file: Output file to save model to + :type output_file: str + """ if not ".keras" in output_file: - output_file+=".keras" - - self._model.save(output_file,) - + output_file += ".keras" + + self._model.save( + output_file, + ) + def load_model(self, input_file: str): """Load model from file @@ -78,7 +87,7 @@ def load_model(self, input_file: str): print(f"Loading model from {input_file}") self._model = tf.keras.models.load_model(input_file) - + def model_predict(self, test_data: pd.DataFrame): """Get model prediction @@ -86,21 +95,23 @@ def model_predict(self, test_data: pd.DataFrame): :type test_data: pd.DataFrame :return: model predction :rtype: NDArray - """ + """ # Hacky but means it's consistent with sci-kit interface scaled_data = self.scale_data(test_data) if self._model is None: return np.zeros(len(test_data)) - + return self._model.predict(scaled_data, verbose=False).T[0] def model_predict_no_scale(self, test_data): # Same as above but specifically for TF, optimised to avoid if statement... return self._model(test_data, training=False) - def evaluate_model(self, predicted_values: Iterable, true_values: Iterable, outfile: str = ""): - + def evaluate_model( + self, predicted_values: Iterable, true_values: Iterable, outfile: str = "" + ): + # CODE TO DO TF SPECIFIC PLOTS GOES HERE - - return super().evaluate_model(predicted_values, true_values, outfile) \ No newline at end of file + + return super().evaluate_model(predicted_values, true_values, outfile) diff --git a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_manual_interface.py b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_manual_interface.py index 94bf203..b0fc441 100644 --- a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_manual_interface.py +++ b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_manual_interface.py @@ -5,36 +5,43 @@ import tensorflow.keras as tfk import keras_tuner as kt + class TfManualInterface(TfInterface): _training_settings = {} - + def set_training_settings(self, kwargs): """Set training settings, needs to be done early for...reasons :param kwargs: training kwargs :type kwargs: kwargs - """ - self._training_settings = kwargs - + """ + self._training_settings = kwargs + def train_model(self): - """train model - """ + """train model""" scaled_data = self.scale_data(self._training_data) scaled_labels = self.scale_labels(self._training_labels) - - lr_schedule = tfk.callbacks.ReduceLROnPlateau(monitor="val_loss", patience=10, factor=0.5, min_lr=1e-8, verbose=1) - stop_early = tfk.callbacks.EarlyStopping(monitor='val_loss', patience=20) - self._model.fit(scaled_data, scaled_labels, **self._training_settings, callbacks=[lr_schedule, stop_early]) + lr_schedule = tfk.callbacks.ReduceLROnPlateau( + monitor="val_loss", patience=10, factor=0.5, min_lr=1e-8, verbose=1 + ) + stop_early = tfk.callbacks.EarlyStopping(monitor="val_loss", patience=20) + + self._model.fit( + scaled_data, + scaled_labels, + **self._training_settings, + callbacks=[lr_schedule, stop_early], + ) class TfManualLayeredInterface(TfManualInterface): __TF_LAYER_IMPLEMENTATIONS = { "dense": tfk.layers.Dense, "dropout": tfk.layers.Dropout, - "batchnorm": tfk.layers.BatchNormalization + "batchnorm": tfk.layers.BatchNormalization, } - + _layers = [] def add_layer(self, layer_id: str, layer_args: dict): @@ -45,20 +52,16 @@ def add_layer(self, layer_id: str, layer_args: dict): :param layer_args: kwargs for layer :type layer_args: dict :raises ValueError: Layer type not implemented in __TF_LAYER_IMPLEMENTATIONS yet - """ + """ if layer_id not in self.__TF_LAYER_IMPLEMENTATIONS.keys(): raise ValueError(f"{layer_id} not implemented yet!") if layer_args.get("kernel_regularizer", False): # Hacky, swaps string value of regularliser for proper one - layer_args["kernel_regularizer"] = tfk.regularizers.L2(layer_args["kernel_regularizer"]) - - self._layers.append(self.__TF_LAYER_IMPLEMENTATIONS[layer_id.lower()](**layer_args)) - - - - - - - + layer_args["kernel_regularizer"] = tfk.regularizers.L2( + layer_args["kernel_regularizer"] + ) + self._layers.append( + self.__TF_LAYER_IMPLEMENTATIONS[layer_id.lower()](**layer_args) + ) diff --git a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_normalizing_flow_model.py b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_normalizing_flow_model.py index dc09714..469773c 100644 --- a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_normalizing_flow_model.py +++ b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_normalizing_flow_model.py @@ -1,4 +1,6 @@ -from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import TfManualInterface +from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import ( + TfManualInterface, +) import tensorflow_probability as tfp import tensorflow.keras as tfk @@ -17,7 +19,7 @@ # def __init__(self, dist): # super().__init__() # self._layer_dist = dist.log_prob - + # def call(self, x): # return self._layer_dist(x) @@ -26,48 +28,55 @@ class NormalizingFlow: def __init__(self, hidden_units: List[int], n_bijectors: int, input_dim: int): self._hidden_units = hidden_units self._n_bijectors = n_bijectors - self._input_dim=input_dim - + self._input_dim = input_dim + # Create an autoregressive bijector def _create_autoregressive_bijector(self): return tfb.MaskedAutoregressiveFlow( shift_and_log_scale_fn=tfb.AutoregressiveNetwork( params=2, # Shift and log scale hidden_units=self._hidden_units, - input_shape=(None,self._input_dim), + input_shape=(None, self._input_dim), ) ) - + def _create_normalizing_flow(self): - bijectors = [self._create_autoregressive_bijector() for _ in range(self._n_bijectors)] + bijectors = [ + self._create_autoregressive_bijector() for _ in range(self._n_bijectors) + ] # Add a final bijector for numerical stability (optional) bijectors.append(tfb.Permute(permutation=list(range(self._input_dim))[::-1])) return tfb.Chain(bijectors) def __call__(self): flow_bijector = self._create_normalizing_flow() - base_distribution = tfd.MultivariateNormalDiag(loc=tf.zeros(self._input_dim), scale_diag=tf.ones(self._input_dim)) + base_distribution = tfd.MultivariateNormalDiag( + loc=tf.zeros(self._input_dim), scale_diag=tf.ones(self._input_dim) + ) # Grab the distribution - return tfd.TransformedDistribution(distribution=base_distribution, bijector=flow_bijector) - - + return tfd.TransformedDistribution( + distribution=base_distribution, bijector=flow_bijector + ) class TfNormalizingFlowModel(TfManualInterface): def build_model(self, **kwargs): - input_dim = self.chain.ndim-1 - self._model = NormalizingFlow(kwargs.get("hidden_units", [100]), kwargs.get("n_bijectors", 1), input_dim)() + input_dim = self.chain.ndim - 1 + self._model = NormalizingFlow( + kwargs.get("hidden_units", [100]), kwargs.get("n_bijectors", 1), input_dim + )() self._optimizer = tfk.optimizers.Adam(kwargs.get("learning_rate", 1e-3)) def nll_loss(self, features): return -tf.reduce_mean(self._model.log_prob(features)) - def train_model(self): epochs = self._training_settings["epochs"] batch_size = self._training_settings["batch_size"] - - scaled_data = tf.data.Dataset.from_tensor_slices(self.scale_data(self._training_data)).batch(batch_size) + + scaled_data = tf.data.Dataset.from_tensor_slices( + self.scale_data(self._training_data) + ).batch(batch_size) for epoch in tqdm(range(epochs)): epoch_loss = 0 @@ -75,7 +84,9 @@ def train_model(self): with tf.GradientTape() as tape: loss = self.nll_loss(batch) grads = tape.gradient(loss, self._model.trainable_variables) - self._optimizer.apply_gradients(zip(grads, self._model.trainable_variables)) + self._optimizer.apply_gradients( + zip(grads, self._model.trainable_variables) + ) epoch_loss += loss.numpy() print(f"Epoch {epoch+1}, Loss: {epoch_loss / len(scaled_data)}") @@ -84,9 +95,16 @@ def test_model(self): surrogate_samples = self._model.sample(10000).numpy() with PdfPages("likelihood_free_inference.pdf") as pdf: - for i in range(self._chain.ndim-1): - sns.histplot(surrogate_samples[:, i], label="Surrogate", color="blue", fill=False) - sns.histplot(self._chain.ttree_array[:, i], label="Test Data", color="blue", fill=False) + for i in range(self._chain.ndim - 1): + sns.histplot( + surrogate_samples[:, i], label="Surrogate", color="blue", fill=False + ) + sns.histplot( + self._chain.ttree_array[:, i], + label="Test Data", + color="blue", + fill=False, + ) plt.legend() plt.xlabel(f"{self._chain.plot_branches[i]}") plt.ylabel("Density") diff --git a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_residual_model.py b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_residual_model.py index 3e537b6..257c5d6 100644 --- a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_residual_model.py +++ b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_residual_model.py @@ -1,10 +1,14 @@ -from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import TfManualLayeredInterface +from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import ( + TfManualLayeredInterface, +) import tensorflow.keras as tfk class TfResidualModel(TfManualLayeredInterface): def build_model(self, **kwargs): - input_shape = self.training_data.shape[1:] # Assuming shape is (batch_size, features) + input_shape = self.training_data.shape[ + 1: + ] # Assuming shape is (batch_size, features) network_input = tfk.layers.Input(shape=input_shape) # Initial layer @@ -23,9 +27,12 @@ def build_model(self, **kwargs): # Define and compile the model self._model = tfk.Model(inputs=network_input, outputs=x) - optimizer = tfk.optimizers.AdamW(learning_rate=kwargs.get("learning_rate", 1e-5), - weight_decay=1e-4, clipnorm=1.0) + optimizer = tfk.optimizers.AdamW( + learning_rate=kwargs.get("learning_rate", 1e-5), + weight_decay=1e-4, + clipnorm=1.0, + ) _ = kwargs.pop("learning_rate", None) - self._model.compile(optimizer=optimizer, **kwargs) \ No newline at end of file + self._model.compile(optimizer=optimizer, **kwargs) diff --git a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_sequential_model.py b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_sequential_model.py index cbab8e2..2bae78e 100644 --- a/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_sequential_model.py +++ b/src/MaCh3PythonUtils/machine_learning/tensorflow/tf_sequential_model.py @@ -1,9 +1,12 @@ -from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import TfManualLayeredInterface +from MaCh3PythonUtils.machine_learning.tensorflow.tf_manual_interface import ( + TfManualLayeredInterface, +) import tensorflow.keras as tfk + class TfSequentialModel(TfManualLayeredInterface): - + def build_model(self, **kwargs: dict): """Build and compile TF model @@ -20,10 +23,12 @@ def build_model(self, **kwargs: dict): self._model.add(layer) self._model.build() - optimizer = tfk.optimizers.AdamW(learning_rate=kwargs.get("learning_rate", 1e-5), - weight_decay=1e-4, clipnorm=1.0) + optimizer = tfk.optimizers.AdamW( + learning_rate=kwargs.get("learning_rate", 1e-5), + weight_decay=1e-4, + clipnorm=1.0, + ) kwargs.pop("learning_rate", None) - - self._model.compile(**kwargs, optimizer=optimizer) \ No newline at end of file + self._model.compile(**kwargs, optimizer=optimizer)