diff --git a/README.md b/README.md index 456af67..9819c73 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,22 @@ Slicing Problem with Reinforcement Learning for DUNE reconstruction. +## Python Environment + +SliceRL requires the following python packages: + +- python3 (version 3.8) +- numpy +- gym +- matplotlib +- tensorflow (version 2.3) +- [keras-rl2](https://github.com/wau/keras-rl2.git) +- [energyflow](https://github.com/pkomiske/EnergyFlow.git) +- json +- gzip +- argparse +- [hyperopt](https://github.com/hyperopt/hyperopt.git) (optional) + ## Problem Slicing algorithm clusters calorimeter hits (CaloHits) based on the main primary interacting particle. This repository implements a DDPG agent that learns to recursively put CaloHits into their correct slice. The agent tries to predict the slice index for each CaloHit. @@ -27,4 +43,8 @@ The [EnergyFlow](https://github.com/pkomiske/EnergyFlow) Python package is used SliceRL feeds a tuple of `(E,x,z)` coordinates to the EDM function. First parameter is the total energy, second is the sum of `x` coordinates of all the CaloHits in the slice, the last is the sum of coordinates in the specific U/V/W plane considered. -Note: For EMD to be a distance metric, the `R` parameter has to be greater than half of the maximum distance between CaloHits. For protoDUNE we have that parameter must be greater than about 600 millimeters for the (x,U/V/W) planes. +Note: For EMD to be a distance metric, the `R` parameter has to be greater than half of the maximum distance between CaloHits. For protoDUNE we have that parameter must be greater than about 600 millimeters for the (x,U/V/W) planes (0.6 in the code). + +## Notes + +Input lenghts are expressed in millimeters, while energies in ADCs. In order to have smaller number and gradients involved in the slicing computations, rescaling factors of 1e3 and 1e2 are employed for lenghts and energies respectively. diff --git a/runcards/default_dense.json b/runcards/default_dense.json new file mode 100644 index 0000000..a3798d9 --- /dev/null +++ b/runcards/default_dense.json @@ -0,0 +1,29 @@ +{ + "slicerl_env": { + "fn": "test_data.csv.gz", + "nev": 5, + "width": 1.0, + "min_hits": 200, + "reward": "cauchy" + }, + "rl_agent": { + "learning_rate": 1e-4, + "nstep": 100000, + "actor": { + "architecture": "Dense", + "dropout": 0.05, + "nb_layers": 3, + "nb_units": 100 + }, + "critic":{ + "architecture": "Dense", + "dropout": 0.05, + "nb_layers": 3, + "nb_units": 100 + }, + "optimizer": "Adam" + }, + "test": { + "fn": "test_data.csv.gz" + } +} diff --git a/runcards/test.json b/runcards/test.json new file mode 100644 index 0000000..4bdf949 --- /dev/null +++ b/runcards/test.json @@ -0,0 +1,30 @@ +{ + "slicerl_env": { + "fn": "test_data_2GeV.csv.gz", + "nev": 5, + "width": 1.0, + "min_hits": 250, + "penalty": 1, + "k": 10, + "reward": "cauchy" + }, + "rl_agent": { + "learning_rate": 1e-4, + "nstep": 10000, + "memory": 5, + "actor": { + "dropout": 0.01, + "nb_layers": 3, + "nb_units": 100 + }, + "critic":{ + "dropout": 0.05, + "nb_layers": 4, + "nb_units": 100 + }, + "optimizer": "Adam" + }, + "test": { + "fn": "test_data_2GeV.csv.gz" + } +} diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..7d0b001 --- /dev/null +++ b/setup.py @@ -0,0 +1,35 @@ +# This file is part of SliceRL by M. Rossi + +from __future__ import print_function +import sys +from setuptools import setup, find_packages + +if sys.version_info < (3,6): + print("SliceRL requires Python 3.6 or later", file=sys.stderr) + sys.exit(1) + +with open('README.md') as f: + long_desc = f.read() + +setup(name= "slicerl", + version = '1.0.0', + description = "Slicing Problem with Reinforcement Learning for DUNE reconstruction", + author = "M. Rossi", + author_email = "marco.rossi@cern.ch", + url="https://github.com/marcorossi5/SlicingRL.git", + long_description = long_desc, + entry_points = {'console_scripts': + ['slicerl = slicerl.scripts.slicerl:main']}, + package_dir = {'': 'src'}, + packages = find_packages('src'), + zip_safe = False, + classifiers=[ + 'Operating System :: Unix', + 'Programming Language :: Python', + 'Programming Language :: Python :: 3', + 'Programming Language :: Python :: 3.5', + 'Topic :: Scientific/Engineering', + 'Topic :: Scientific/Engineering :: Physics', + ], + ) + diff --git a/src/slicerl/AgentSlicerl.py b/src/slicerl/AgentSlicerl.py new file mode 100644 index 0000000..f241af6 --- /dev/null +++ b/src/slicerl/AgentSlicerl.py @@ -0,0 +1,44 @@ +# This file is part of SliceRL by M. Rossi + +from slicerl.Slicer import ContinuousSlicer +from rl.agents.ddpg import DDPGAgent +from tensorflow.keras.models import model_from_json + +#====================================================================== +class DDPGAgentSlicerl(DDPGAgent): + """DDPG Agent for pile up subtraction""" + + #---------------------------------------------------------------------- + def __init__(self, *args, **kwargs): + """Initialize the DQN agent.""" + super(DDPGAgentSlicerl, self).__init__(*args, **kwargs) + + #---------------------------------------------------------------------- + def save(self, filepath, overwrite=False, include_optimizer=True): + """Save the model to file.""" + self.model.save(filepath, overwrite=overwrite, + include_optimizer=include_optimizer) + + #---------------------------------------------------------------------- + def load_model(self, filepath, custom_objects=None, compile=True): + """Load model from file""" + self.model = load_model(filepath, custom_objects=custom_objects, + compile=compile) + self.update_target_model_hard() + + #---------------------------------------------------------------------- + def load_with_json(self, jsonfile, weightfile): + """ + Load model from a json file with the architecture, and an h5 file with weights. + """ + # read architecture card + with open(jsonfile) as f: + arch = json.load(f) + self.model = model_from_json(arch) + self.model.load_weights(weightfile) + self.update_target_model_hard() + + #---------------------------------------------------------------------- + def slicer(self): + """Return the current subtractor.""" + return ContinuousSlicer(self.actor) diff --git a/src/slicerl/Event.py b/src/slicerl/Event.py new file mode 100644 index 0000000..5af00d9 --- /dev/null +++ b/src/slicerl/Event.py @@ -0,0 +1,208 @@ +# This file is part of SliceRL by M. Rossi + +import sys +import numpy as np +import math +import json, gzip +from collections import namedtuple + +EventTuple = namedtuple("EventTuple", ["E", "x", "z", "cluster_idx", "pndr_idx", + "mc_idx", "slicerl_idx"]) + +#====================================================================== +class Event: + """Event class keeping track of calohits in a 2D plane view.""" + + #---------------------------------------------------------------------- + def __init__(self, calohits, k, min_hits=1): + """ + Parameters + ---------- + - calohits : np.array, shape=(7, num calohits), containing all the + event information. Each calohit is described by: + - energy [ADC/100] + - x coordinate [10^1 m] + - z coordinate [10^1 m] + - input pfo list cluster idx + - pandora slice index + - cheating slice idx (mc truth) + - slicing output (array of minus ones if not loading + inference results) + - min_hits : int, consider only slices with equal or more than min_hits + - Lneigh : neighbourhood radius [10^1 m] + """ + self.k = k + self.calohits = [] + for i,c in enumerate(calohits.T): + # find calohits indices belonging to the same slice + in_slice_m = np.where(calohits[5] == c[5])[0] + + include = len(in_slice_m) >= min_hits + mcstate = np.array( + [calohits[0][in_slice_m].sum(), # cumulative energy + calohits[0][in_slice_m].sum(), # cumulative x + calohits[0][in_slice_m].sum()] # cumulative z + ) + # find the k closest calohits in (x,z) plane + sqdist = (calohits[1] - c[1])**2 + (calohits[2] - c[2])**2 + idx = np.argpartition(sqdist,1+self.k)[:1+self.k] + sorted_idx = idx[np.argsort(sqdist[idx])] + neighbourhood = Graph(calohits[:, sorted_idx], + sorted_idx, + sqdist[sorted_idx]) + + # calohits with status different from -1 or None will never be + # processed. The model focuses on calohits with more calohits + # than min_hits in the cluster, but the graph is allowed to + # consider all the calohits. + status = c[6] if include else -2. + self.calohits.append(CaloHit(c[0], c[1], c[2],c[3], c[4], c[5], + neighbourhood=neighbourhood, + mc_neighbours_m=in_slice_m, mcstate=mcstate, + status=status)) + + #---------------------------------------------------------------------- + def __len__(self): + return len(self.calohits) + + #---------------------------------------------------------------------- + def state(self, i): + """ + Return the observable state of the calohit at index i. + + Returns + ------- + - array of shape=(1+k,6) + """ + if i < len(self.calohits): + return self.calohits[i].neighbourhood.state(self.calohits) + else: + return np.zeros([1+self.k, 6]) + + #---------------------------------------------------------------------- + def calohits_to_array(self): + """ + Takes calohits list into a list of arrays for plot rendering. + + Returns + ------- + - numpy array of shape (7, num calohits). Rows contain in order: + energies, xs, zs, cluster_idx, pndr_idx, cheating_idx, slicerl_idx + + Note + ---- + - energy is measured in ADC + - spatial coordinates x and z are converted in cm + + + """ + rows = [[] for i in range(7)] + for c in self.calohits: + rows[0].append(c.E*100) + rows[1].append(c.x*1000) + rows[2].append(c.z*1000) + rows[3].append(c.clusterIdx) + rows[4].append(c.pndrIdx) + rows[5].append(c.mcIdx) + rows[6].append(c.status) + return np.array(rows) + + #---------------------------------------------------------------------- + def calohits_to_namedtuple(self): + """ + Takes calohits list into an EventTuple object for plot rendering. + + Returns + ------- + - numpy array of shape (7, num calohits). Rows contain in order: + energies, xs, zs, cluster_idx, pndr_idx, cheating_idx, slicerl_idx + + """ + arr = self.calohits_to_array() + return EventTuple(*arr) + + #---------------------------------------------------------------------- + def dump(self, fname): + """ Dump calohits list to fname file """ + rows = self.calohits_to_array() + for row in rows: + np.savetxt(fname, row, fmt='%.5f', newline=',') + fname.write(b"\n") + +#====================================================================== +class CaloHit: + """Keep track of calohit properties.""" + + #---------------------------------------------------------------------- + def __init__(self, E, x, z, clusterIdx=-1, pndrIdx=None, mcIdx=None, neighbourhood=None, + mc_neighbours_m=None, mcstate=None, status=None): + """ + Create a calohit that has E,x,z members as well as + - E : the energy carried by the calohit [10^-2 ADC] + - x : the drift coordinate of the calohit [10^1 m] + - z : the 2D plane coordinate of the calohit [10^1 m] + - clusterIdx : additional information from other pandora algorithms, + -1 if not provided + - pndrIdx : current pandora slicing algorithm output + - mcIdx : truth slice index + - neighbourhood : Graph object describing calohit neighbourhood + - mc_neighbours_m : mask array of calohits from the same slice (mc truth) (TODO: remove this?) + - mc_state : the cumulative (E,x,z) info about the slice the calohit + belongs to (mc truth) + - status : the current slice index the calohit belongs to (changed + by the agent) + """ + self.E = E + self.x = x + self.z = z + self.clusterIdx = clusterIdx + self.pndrIdx = pndrIdx + self.mcIdx = mcIdx + self.neighbourhood = neighbourhood + self.mc_neighbours_m = mc_neighbours_m if mc_neighbours_m is not None else [] + self.mc_state = mcstate + self.status = status + +#====================================================================== +class Graph: + """ Shape calohit neighbourhood into a graph. """ + + #---------------------------------------------------------------------- + def __init__(self, calohits, idxs, sqdist): + """ + Build the Graph object around a specific calohit. + + Parameters + ---------- + - calohits : array of calohits feature in the event, shape=(7, 1+k) + - idxs : array, ordered neighbours indices, shape=(1+k,) + - sqdist : array of squared distances from central node, shape=(1+k,) + + Note + ---- + all arrays are ordered by increasing calohit distance from center in + the (x,z) plane + """ + self.neighbours_idxs = idxs # neighbour calohits indices + self.fv = np.concatenate([[sqdist], + calohits[(0,1,2),:]]) # feature vector + self.cv = calohits[3,:] # cluster idx vector + + #---------------------------------------------------------------------- + def state(self, calohits): + """ + Return the graph node feature vector concatenated with the current value + of Calohit status attribute of calohits in the neighbourhood. + + Parameter + --------- + - calohits : list of Calohit objects in the neighbourhood + + Returns + ------- + array of shape=(1+k, 6) + """ + sv = np.array( [calohits[idx].status for idx in self.neighbours_idxs] ) + return np.concatenate([self.fv, [self.cv/100], [sv/100]]).T + +# TODO: think about normalizing the cluster_idx and status inputs to the actor model \ No newline at end of file diff --git a/src/slicerl/Slicer.py b/src/slicerl/Slicer.py new file mode 100644 index 0000000..62d3206 --- /dev/null +++ b/src/slicerl/Slicer.py @@ -0,0 +1,104 @@ +# This file is part of SliceRL by M. Rossi + +import os +import numpy as np +import math, json +from abc import ABC, abstractmethod +from rl.policy import GreedyQPolicy +from tensorflow.keras.models import model_from_json + +#====================================================================== +class AbstractSlicer(ABC): + """AbstractSlicer class.""" + + #---------------------------------------------------------------------- + def __call__(self, event): + """Apply the subtractor to an event and returned the cleaned event.""" + # TODO: implement + return self._slicing(event) + + #---------------------------------------------------------------------- + @abstractmethod + def _slicing(self, event): + return event + +class ContinuousSlicer(AbstractSlicer): + """ContinuousSlicer class that acts on an event using a discrete action.""" + + #---------------------------------------------------------------------- + def __init__(self, actor=None): + """Initialisation of the subtractor.""" + # load a ddpg instance here + self.actor = actor + self.action_max = 1.0 + self.eps = np.finfo(np.float32).eps + + #---------------------------------------------------------------------- + def _slicing(self, event): + """ + Apply subtractor to an event. Just overwrite calohit status with slice + index, do not keep track of slices cumulative statistics. Delegate + slice checks to diagnostics. + """ + slice_count = 0 + + for i in range(len(event.calohits)): + c = event.calohits[i] + if i == 0: + # first calohit seeds a new slice automatically + c.status = 0 + slice_count += 1 + continue + + state = event.state(i) + action = self.actor.predict_on_batch(np.array([[state]])).flatten() + action = np.clip(action, -self.action_max, self.action_max-self.eps) + + if action[0] >= 0: + idx = math.floor(action*slice_count) + c.status = idx + else: + c.status = slice_count + slice_count += 1 + + return event + + #---------------------------------------------------------------------- + def load_with_json(self, jsonfile, weightfile): + """ + Load model from a json file with the architecture, and an h5 file with weights. + First actor, second critic. + """ + # read actor architecture card + filename, extension = os.path.splitext(jsonfile) + actor_jsonfile = filename + '_actor' + extension + with open(actor_jsonfile) as f: + arch = json.load(f) + self.actor = model_from_json(arch) + + filename, extension = os.path.splitext(weightfile) + actor_weightfile = filename + '_actor' + extension + self.actor.load_weights(actor_weightfile) + + #---------------------------------------------------------------------- + def save(self, filepath, overwrite=False, include_optimizer=True): + """Save the model to file.""" + self.actor.save(filepath, overwrite=overwrite, + include_optimizer=include_optimizer) + + #---------------------------------------------------------------------- + def load_model(self, filepath, custom_objects=None, compile=True): + """Load model from file""" + # TODO: check this + self.actor = load_model(filepath, custom_objects=custom_objects, + compile=compile) + + #---------------------------------------------------------------------- + def save_weights(self, filepath, overwrite=False): + """Save the weights of model to file.""" + self.actor.save_weights(filepath, overwrite=overwrite) + + #---------------------------------------------------------------------- + def load_weights(self, filepath): + """Load weights of model from file""" + self.actor.load_weights(filepath) diff --git a/src/slicerl/SlicerlEnv.py b/src/slicerl/SlicerlEnv.py new file mode 100644 index 0000000..2cf7fb1 --- /dev/null +++ b/src/slicerl/SlicerlEnv.py @@ -0,0 +1,208 @@ +# This file is part of SliceRL by M. Rossi + +import random, math, pprint +from slicerl.read_data import Events +from slicerl.Event import Event +from gym import spaces, Env +from gym.utils import seeding +import numpy as np +from slicerl.tools import quality_metric +from copy import deepcopy + +#====================================================================== +class SlicerlEnv(Env): + """Class defining a gym environment for the slicing algorithm.""" + #---------------------------------------------------------------------- + def __init__(self, hps, low, high): + """ + Initialisation of the environment using a dictionary with hyperparameters + and a lower and upper bound for the observable state. + + The hps dictionary should have the following entries: + - fn: filename of data set + - nev: number of events (-1 for all) + - reward: type of reward function (cauchy, gaussian, ...) + """ + # add small penalty for not clustering + self.penalty = hps['penalty'] + + # read in the events + self.k = hps['k'] + reader = Events(hps['fn'], hps['nev'], k=hps['k'], min_hits=hps['min_hits']) + self.events = reader.values() + + # set up the count parameters and initial state + self.slice_count = 0 + self.slices = [] + + # set up action and observation space + self.action_space = spaces.Discrete(2) + self.observation_space = spaces.Box(low, high, dtype=np.float32) + + # set up some internal parameters + self.seed() + self.viewer = None + + # set the reward function + self.width = hps['width'] + if hps['reward']=='cauchy': + self.__reward=self.__reward_Cauchy + elif hps['reward']=='gaussian': + self.__reward=self.__reward_Gaussian + elif hps['reward']=='exponential': + self.__reward=self.__reward_Exponential + elif hps['reward']=='inverse': + self.__reward=self.__reward_Inverse + else: + raise ValueError('Invalid reward: %s'%hps['reward']) + + self.description= '%s with\n%s' % (self.__class__.__name__,pprint.pformat(hps)) + print('Setting up %s' % self.description) + + #---------------------------------------------------------------------- + def reset_current_event(self): + """Reset the current event and insert the first calohit.""" + self.event = deepcopy(random.choice(self.events)) + self.index = -1 + self.slice_count = 0 + self.slices = [] + self.set_next_node() + + # overwrite calohit status with slice index + + c = self.event.calohits[self.index] + c.status = self.slice_count + + # add calohit (E,x,z) to cumulative slice state + calohit_state = self.state[0,1:4] + self.slices.append( calohit_state ) + self.slice_count += 1 + + self.set_next_node() + + #---------------------------------------------------------------------- + def set_next_node(self, is_reset=False): + """Set the current particle using the event list.""" + # print('setting node up') + while self.index < len(self.event): + self.index += 1 + if self.index < len(self.event): + if (self.event.calohits[self.index].status is None)\ + or (self.event.calohits[self.index].status == -1): + break + self.state = self.event.state(self.index) + + #---------------------------------------------------------------------- + def __reward_Cauchy(self, x): + """A cauchy reward function.""" + return 1.0/(math.pi*(1.0 + (x*x))) + + #---------------------------------------------------------------------- + def __reward_Gaussian(self, x): + """A gaussian reward function.""" + return np.exp(-x*x/2.0) + + #---------------------------------------------------------------------- + def __reward_Exponential(self, x): + """A negative exponential reward function.""" + return np.exp(-x) + + #---------------------------------------------------------------------- + def __reward_Inverse(self, x): + """An inverse reward function.""" + return min(1.0, 1.0/(x + 0.5)) + + #---------------------------------------------------------------------- + def reward(self, slice_state, mc_state): + """Full reward function.""" + x = quality_metric(slice_state, mc_state) + return self.__reward(x/self.width) + + #---------------------------------------------------------------------- + def seed(self, seed=None): + """Initialize the seed.""" + self.np_random, seed = seeding.np_random(seed) + return [seed] + + #---------------------------------------------------------------------- + def reset(self): + """Reset the event to a new randomly selected one.""" + self.reset_current_event() + return self.state + + #---------------------------------------------------------------------- + def render(self, mode='human'): + """Add potential intermediary output here""" + pass + + #---------------------------------------------------------------------- + def close(self): + if self.viewer: self.viewer.close() + +#====================================================================== +class SlicerlEnvContinuous(SlicerlEnv): + """Class defining a gym environment for the continuous slicing algorithm.""" + #---------------------------------------------------------------------- + def __init__(self, *args, **kwargs): + super(SlicerlEnvContinuous, self).__init__(*args, **kwargs) + self.eps = np.finfo(np.float32).eps + self.action_max = 1.0 + self.action_space = spaces.Box(low=-self.action_max, high=self.action_max, + shape=(1,), dtype=np.float32) + + #---------------------------------------------------------------------- + def step(self, action): + """ + Perform a step using the current calohit in the event, deciding which + slice to add it to. + """ + action = np.clip(action, -self.action_max, self.action_max-self.eps) + assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action)) + c = self.event.calohits[self.index] + mc_state = c.mc_state + calohit_state = self.state[0,1:4] + + aggregate = action[0] >= 0 + if aggregate: + # select the correct existing slice to put the calohit in + idx = math.floor(action[0]*self.slice_count) + + # overwrite calohit status attribute with slice index + c.status = idx + + # add calohit (E,x,z) to cumulative slice state + self.slices[idx] += calohit_state + slice_state = self.slices[idx] + else: + # form a new slice + # overwrite calohit status with slice index + c.status = self.slice_count + + # add calohit (E,x,z) to cumulative slice state + self.slices.append( calohit_state ) + slice_state = calohit_state + self.slice_count += 1 + + # calculate a reward + # penalize wrongly seeding new slice and when aggregating to small slice + large_cluster = np.count_nonzero(c.mc_neighbours_m) >= 5 + penalty = self.penalty / 1e3 * (not aggregate) + extra_reward = self.penalty / 1e3 * (large_cluster == aggregate) + extra_penalty = self.penalty / 5e1 * (large_cluster and not aggregate) + # penalize when aggregating if neighbours < 5 + # penalize when seeding new slide if neighbours > 5 + reward = self.reward(slice_state, mc_state) + extra_reward - penalty - extra_penalty + + # move to the next node in the clustering sequence + self.set_next_node() + + # if we are at the end of the declustering list, then we are done for this event. + done = bool(self.index >= len(self.event)) + + # return the state, reward, and status + return self.state, reward, done, {} + +""" +TODO: split the observation space in 3-dim Box plut +TODO: implement the env.step function +""" \ No newline at end of file diff --git a/src/slicerl/__init__.py b/src/slicerl/__init__.py new file mode 100644 index 0000000..1bc2631 --- /dev/null +++ b/src/slicerl/__init__.py @@ -0,0 +1 @@ +#We don't want to import stuff here that could slow down the import times diff --git a/src/slicerl/diagnostics.py b/src/slicerl/diagnostics.py new file mode 100644 index 0000000..7adca69 --- /dev/null +++ b/src/slicerl/diagnostics.py @@ -0,0 +1,282 @@ +# This file is part of SliceRL by M. Rossi +import os +from tensorflow.keras.utils import Progbar +from slicerl.Event import Event +from slicerl.tools import quality_metric +from matplotlib import pyplot as plt +import numpy as np +import math + +#---------------------------------------------------------------------- +colors = [ + "lightsalmon", + "orange", + "springgreen", + "fuchsia", + "lime", + "lightcoral", + "pink", + "darkseagreen", + "gold", + "red", + "deepskyblue", + "lightgreen", + "coral", + "aqua", + "lightgreen", + "mediumaquamarine" +] +l = len(colors) + +#---------------------------------------------------------------------- +def print_stats(name, data, mass_ref, output_folder='./'): + """Print statistics on the mass distribution.""" + r_plain = np.array(data)-mass_ref + m = np.median(r_plain) + a = np.mean(r_plain) + s = np.std(r_plain) + fname = f"{output_folder}/diagnostics.txt" + with open(fname,'a+') as f: + print(f"{name:<25}:\tmedian-diff {m:.2f}\tavg-diff {a:.2f}\tstd-diff {s:.2f}", + file=f) + +#---------------------------------------------------------------------- +def plot_multiplicity(events, output_folder='./', loaddir=None): + """Plot the slice multiplicity distribution and output some statistics.""" + if loaddir is not None: + fname = '%s/nmc.npy' % loaddir + nmc = np.load(fname) + fname = '%s/nddpg.npy' % loaddir + nddpg = np.load(fname) + else: + nmc = np.array( [int(event.mc_idx.max()) + 1 for event in events] ) + nddpg = np.array( [int(event.slicerl_idx.max()) + 1 for event in events] ) + + bins = np.linspace(0, 200, 201) + hnmc, _ = np.histogram(nmc, bins=bins) + hnddpg, _ = np.histogram(nddpg, bins=bins) + + plt.rcParams.update({'font.size': 20}) + plt.figure(figsize=(18,14)) + + plt.hist(bins[:-1], bins, weights=hnmc, histtype='step', color='blue', label='mc') + plt.hist(bins[:-1], bins, weights=hnddpg, histtype='step', color='red', label='ddpg') + + plt.xlabel("multiplicity", loc='right') + plt.xlim((bins[0], bins[-1])) + plt.legend() + fname = f"{output_folder}/multiplicity.pdf" + plt.savefig(fname, bbox_inches='tight') + + print_stats('Slice multiplicity', nddpg , nmc, output_folder=output_folder) + resultsdir = '%s/results' % output_folder + np.save(f"{resultsdir}/nmc.npy", nmc) + np.save(f"{resultsdir}/nddpg.npy", nddpg) + +#---------------------------------------------------------------------- +def plot_EMD(events, events_obj, output_folder='./', loaddir=None): + """Plot the energy movers distance between slices and output some statistics.""" + if loaddir is not None: + fname = '%s/emdddpg.npy' % loaddir + emdddpg = np.load(fname) + else: + emdddpg = [] + for i, event in enumerate(events): + num_clusters = event.slicerl_idx.max().astype(np.int32) + 1 + for idx in range(num_clusters): + m = event.slicerl_idx == idx + Es = event.E[m].sum() + xs = event.x[m].sum() + zs = event.z[m].sum() + slice_state = np.array([Es, xs, zs]) + # get the index of the first calohit in the slice + cidx = np.argwhere(m)[0,0] + mc_state = events_obj[i].calohits[cidx].mc_state + emdddpg.append(quality_metric(slice_state, mc_state)) + emdddpg = np.array(emdddpg) + + bins = np.linspace(0, 200, 201) + hemdddpg, _ = np.histogram(emdddpg, bins=bins) + + plt.rcParams.update({'font.size': 20}) + plt.figure(figsize=(18,14)) + + plt.hist(bins[:-1], bins, weights=hemdddpg, histtype='step', color='red') + + plt.xlabel("emd", loc='right') + plt.xlim((bins[0], bins[-1])) + fname = f"{output_folder}/emd.pdf" + plt.savefig(fname, bbox_inches='tight') + + print_stats('Slice EMD', emdddpg, 0, output_folder=output_folder) + resultsdir = '%s/results' % output_folder + np.save(f"{resultsdir}/emdddpg.npy", emdddpg) + +#---------------------------------------------------------------------- +def plot_slice_size(events, output_folder='./', loaddir=None): + """Plot the slice size distribution and output some statistics.""" + bins = np.linspace(0, 100, 101) + + if loaddir is not None: + fname = '%s/smc.npy' % loaddir + smc = np.load(fname) + fname = '%s/sddpg.npy' % loaddir + sddpg = np.load(fname) + else: + binc_mc = [np.bincount(event.mc_idx.astype(np.int32)[event.mc_idx >= 0]) for event in events] + binc_ddpg = [np.bincount(event.slicerl_idx.astype(np.int32)) for event in events] + smc = sum( [np.histogram(bc, bins=bins)[0] for bc in binc_mc] ) + sddpg = sum( [np.histogram(bc, bins=bins)[0] for bc in binc_ddpg] ) + + plt.rcParams.update({'font.size': 20}) + plt.figure(figsize=(18,14)) + + plt.hist(bins[:-1], bins, weights=smc, histtype='step', color='blue', label='mc') + plt.hist(bins[:-1], bins, weights=sddpg, histtype='step', color='red', label='ddpg') + + plt.xlabel("size", loc='right') + plt.xlim((bins[0], bins[-1])) + + plt.legend() + fname = f"{output_folder}/slice_size.pdf" + plt.savefig(fname, bbox_inches='tight') + + print_stats('Slice multiplicity', sddpg , smc, output_folder=output_folder) + resultsdir = '%s/results' % output_folder + np.save(f"{resultsdir}/smc.npy", smc) + np.save(f"{resultsdir}/sddpg.npy", sddpg) + +#---------------------------------------------------------------------- +def plot_plane_view(events, output_folder='./', loaddir=None): + event_arr = events[0] + + fig = plt.figure(figsize=(18*2,14)) + ax = fig.add_subplot(121) + ax.set_title(f"Slicing Algorithm Output, 2D plane view") + ax.set_xlabel("x [mm]") + ax.set_ylabel("z [mm]") + num_clusters = int(event_arr.slicerl_idx.max()) + 1 + sort_fn = lambda x: np.count_nonzero( event_arr.slicerl_idx == x ) + sorted_indices = sorted( range(num_clusters), key=sort_fn, reverse=True) + # sort all the cluster with greater number of hits + print(f"Plotting {num_clusters} Slices in event over {event_arr.slicerl_idx.shape[0]} particles") + for index in sorted_indices[:min(len(sorted_indices), 200)]: + m = event_arr.slicerl_idx == index + ax.scatter(event_arr.x[m], event_arr.z[m], marker='.', color=colors[index%l]) + ax.set_box_aspect(1) + + ax = fig.add_subplot(122) + ax.set_title(f"Cheating Algorithm Truths, 2D plane view") + ax.set_xlabel("x [mm]") + ax.set_ylabel("z [mm]") + num_clusters = int(event_arr.mc_idx.max()) + 1 + print(f"Plotting {num_clusters} true Slices in event") + for index in range(num_clusters): + m = event_arr.mc_idx == index + ax.scatter(event_arr.x[m], event_arr.z[m], marker='.', color=colors[index%l]) + ax.set_box_aspect(1) + fname = f"{output_folder}/pview.pdf" + plt.savefig(fname, bbox_inches='tight') + +#---------------------------------------------------------------------- +def plot_ROC(scores, output_folder='./', loaddir=None): + """Plot the ROC curve for particle classification and output some statistics.""" + # particle: list of np.arrays of shape=(num particles, 2) for all events + if loaddir is not None: + fname = '%s/confusion.npy' % loaddir + tp, fp, fn, tn = np.load(fname) + else: + tp, fp, fn, tn = confusion_matrix_per_event(scores) + efficiency = tp / (tp + fn) + irejection = fp / (tn + fp) + + sqrtn = math.sqrt(len(tp)) + tp_avg = tp.mean() * 100 + tp_unc = tp.std() / sqrtn * 100 + fn_avg = fn.mean() * 100 + fn_unc = fn.std() / sqrtn * 100 + fp_avg = fp.mean() * 100 + fp_unc = fp.std() / sqrtn * 100 + tn_avg = tn.mean() * 100 + tn_unc = tn.std() / sqrtn * 100 + PU_unc = np.sqrt(tp_unc**2 + fn_unc**2) + LV_unc = np.sqrt(tn_unc**2 + fp_unc**2) + + print(f"Dataset balancing: \t PU particles: {tp_avg+fn_avg:.2f}+-{PU_unc:.2f} % \t LV particles: {tn_avg+fp_avg:.2f}+-{LV_unc:.2f} % ") + print("Average confusion matrix:") + print( "\t\ true | | |") + print( "\t ---- | PU | LV |") + print( "\t pred \| | |") + print( "\t---------------------------") + print(f"\t PU |{tp_avg:>5.2f} % |{fp_avg:>5.2f} % |") + print( "\t---------------------------") + print(f"\t LV |{fn_avg:>5.2f} % |{tn_avg:>5.2f} % |") + print( "\t---------------------------") + + plt.rcParams.update({'font.size': 20}) + plt.figure(figsize=(18,14)) + plt.scatter(efficiency, irejection, marker='.', color="blue", label="DQN-Subtracting") + + plt.xlabel("Efficiency $\epsilon$", loc='right') + plt.ylabel("$1/R$", loc="top") + plt.xlim((0,1)) + plt.ylim((0,1)) + plt.legend() + fname = f"{output_folder}/ROC.pdf" + plt.savefig(fname, bbox_inches='tight') + + print_stats('Efficiency' , efficiency , 0, output_folder=output_folder) + print_stats('1/Rejection', irejection , 0, output_folder=output_folder) + resultsdir = '%s/results' % output_folder + np.save(f"{resultsdir}/confusion.npy", np.stack([tp, fp, fn, tn])) + +#---------------------------------------------------------------------- +def inference(slicer, events): + """ + Slice calohits from a list of Events objects. Returns the list of + processed Events. + + Parameters + ---------- + slicer: Slicer object + + Returns + ------- + The list of subtracted Events. + """ + progbar = Progbar(len(events)) + for i, event in enumerate(events): + slicer(event) + progbar.update(i+1) + return events + +#---------------------------------------------------------------------- +def make_plots(events_obj, plotdir): + """ + Make diagnostics plots from a list of subtracted events. + + Parameters + ---------- + - events_obj: list, list of sliced Event objects + - plotdir: str, plots output folder + """ + events = [event.calohits_to_namedtuple() for event in events_obj] + plot_multiplicity(events, plotdir) + plot_slice_size(events, plotdir) + plot_plane_view(events, plotdir) + plot_EMD(events, events_obj, plotdir) + +#---------------------------------------------------------------------- +def load_and_dump_plots(plotdir, loaddir): + """ + Make diagnostics plots from plot data contained in loaddir. + + Parameters + ---------- + - plotdir: str, plots output folder + - loaddir: str, directory where to load plot data from + """ + plot_multiplicity(None, plotdir, loaddir) + plot_slice_size(None, plotdir, loaddir) + plot_plane_view(None, plotdir, loaddir) + plot_EMD(None, None, plotdir, loaddir) diff --git a/src/slicerl/keras_to_cpp.py b/src/slicerl/keras_to_cpp.py new file mode 100644 index 0000000..e73d196 --- /dev/null +++ b/src/slicerl/keras_to_cpp.py @@ -0,0 +1,44 @@ +# This file is part of SliceRL by M. Rossi + +#---------------------------------------------------------------------- +def check_model(hps): + """Check that the model defined is portable to cpp.""" + if hps['dropout']>0.0 or hps['architecture']=='LSTM': + raise ValueError("keras_to_cpp: Only Dense layers without Dropout are supported.") + +#---------------------------------------------------------------------- +# adapted from dump_to_simple_cpp.py from the keras2cpp library +def keras_to_cpp(model, layerdic, cpp_fn): + """Convert keras to cpp readable file.""" + with open(cpp_fn, 'w') as fout: + fout.write('layers ' + str(len(model.layers)) + '\n') + + layers = [] + for ind, l in enumerate(layerdic): + fout.write('layer ' + str(ind) + ' ' + l['class_name'] + '\n') + layers += [l['class_name']] + + if l['class_name'] == 'Conv2D': + W = model.layers[ind].get_weights()[0] + fout.write(str(W.shape[0]) + ' ' + str(W.shape[1]) + ' ' + str(W.shape[2]) + ' ' + str(W.shape[3]) + ' ' + l['config']['padding'] + '\n') + for i in range(W.shape[0]): + for j in range(W.shape[1]): + for k in range(W.shape[2]): + fout.write(str(W[i,j,k]) + '\n') + fout.write(str(model.layers[ind].get_weights()[1]) + '\n') + + if l['class_name'] == 'Activation': + fout.write(l['config']['activation'] + '\n') + + if l['class_name'] == 'MaxPooling2D': + fout.write(str(l['config']['pool_size'][0]) + ' ' + str(l['config']['pool_size'][1]) + '\n') + #if l['class_name'] == 'Flatten': + # print(l['config']['name']) + + if l['class_name'] == 'Dense': + #fout.write(str(l['config']['output_dim']) + '\n') + W = model.layers[ind].get_weights()[0] + fout.write(str(W.shape[0]) + ' ' + str(W.shape[1]) + '\n') + for w in W: + fout.write(str(w) + '\n') + fout.write(str(model.layers[ind].get_weights()[1]) + '\n') diff --git a/src/slicerl/models.py b/src/slicerl/models.py new file mode 100644 index 0000000..a653dc1 --- /dev/null +++ b/src/slicerl/models.py @@ -0,0 +1,227 @@ +# This file is part of SliceRL by M. Rossi + +from slicerl.SlicerlEnv import SlicerlEnvContinuous +import numpy as np + +from slicerl.tools import get_window_width, mass +from slicerl.AgentSlicerl import DDPGAgentSlicerl +from slicerl.read_data import Events +from rl.memory import SequentialMemory +from rl.random import OrnsteinUhlenbeckProcess + +import tensorflow as tf +from tensorflow.keras.models import Sequential, Model +from tensorflow.keras.layers import Dense, Activation, Flatten, Dropout, Input, Concatenate, Reshape +from tensorflow.keras.optimizers import Adam, SGD, RMSprop, Adagrad +import tensorflow.keras.backend as K +from tensorflow.keras.callbacks import TensorBoard + +from hyperopt import STATUS_OK +from time import time + +import pprint, json + +#---------------------------------------------------------------------- +def Block(hps, name, activation): + model = Sequential(name=name) + for i in range(hps['nb_layers']): + model.add(Dense(hps['nb_units'])) + model.add(Activation('relu')) + if hps['dropout']>0.0: + model.add(Dropout(hps['dropout'])) + model.add(Dense(1)) + model.add(Activation(activation)) + return model + +#---------------------------------------------------------------------- +def build_actor_model(hps, input_dim, k): + """Construct the actor model used by the DDPG. Outputs the action""" + input_state = Input(shape=(1,) + input_dim, name='actor_input') + graphBlock = Block(hps, 'GraphBlock', 'linear') + aggregatorBlock = Block(hps, 'AggregatorBlock', 'linear') + finalBlock = Block(hps, 'FinalBlock', 'tanh') + + reshaped_state = Reshape(input_dim)(input_state) + c, notc = tf.split(reshaped_state, [1,k], -2) + c = Reshape((6,))(c) + + neigh_state = graphBlock(notc) + neigh_state = Reshape((k,))(neigh_state) + aggr_state = aggregatorBlock(neigh_state) + + cat_state = Concatenate(axis=-1)([c, aggr_state]) + out = finalBlock(cat_state) + + model = Model(inputs=input_state, outputs=out, name="Actor") + model.summary() + return model + +#---------------------------------------------------------------------- +def build_critic_model(hps, input_dim, k): + """ + Construct the critic model used by the DDPG. Judge the goodness of the + predicted action + """ + action_input = Input(shape=(1,), name='action_input') + obs_input = Input(shape=(1,)+input_dim, name='observation_input') + criticBlock = Block(hps, 'critic', 'linear') + + c, _ = tf.split(obs_input, [1,k], -2) + flattened_obs = Flatten()(c) + x = Concatenate()([action_input, flattened_obs]) + x = criticBlock(x) + + model = Model(inputs=[action_input, obs_input], outputs=x, name="Critic") + model.summary() + return model, action_input + +#---------------------------------------------------------------------- +def build_ddpg(hps, input_dim, k): + """Create a DDPG agent to be used on pandora inputs.""" + + print('[+] Constructing DDPG agent, model setup:') + pprint.pprint(hps) + + # set up the agent + K.clear_session() + actor_model = build_actor_model(hps['actor'], input_dim, k) + critic_input_dim = input_dim + critic_model, action_input = build_critic_model(hps['critic'], critic_input_dim, k) + + memory = SequentialMemory(limit=500000, window_length=1) + random_process = OrnsteinUhlenbeckProcess(size=1, theta=.15, mu=0., sigma=.3) + agent = DDPGAgentSlicerl(actor=actor_model, critic=critic_model, + critic_action_input=action_input, nb_actions=1, + memory=memory, nb_steps_warmup_actor=100, + nb_steps_warmup_critic=100, target_model_update=1e-2, + random_process=random_process) + + if hps['optimizer'] == 'Adam': + opt = Adam(lr=hps['learning_rate']) + elif hps['optimizer'] == 'SGD': + opt = SGD(lr=hps['learning_rate']) + elif hps['optimizer'] == 'RMSprop': + opt = RMSprop(lr=hps['learning_rate']) + elif hps['optimizer'] == 'Adagrad': + opt = Adagrad(lr=hps['learning_rate']) + + agent.compile(opt, metrics=['mae']) + + return agent + +#---------------------------------------------------------------------- +def load_runcard(runcard): + """Read in a runcard json file and set up dimensions correctly.""" + with open(runcard,'r') as f: + res = json.load(f) + # if there is a state_dim variable, set up LundCoordinates accordingly + # unless we are doing a scan (in which case it needs to be done later) + env_setup = res.get("slicerl_env") + return res + +#---------------------------------------------------------------------- +def loss_calc(dqn, fn, nev): + pass + # reader = Events(fn, nev) + # subtractor = dqn.subtractor() + # self.jet_pairs_list = reader.values() + # for jet_pairs in reader: + # if len(jet_pairs)==0: + # continue + # events = Event(jet_pairs) + # for event in events: + # if + # reader_sig = Jets(fn_sig, nev) # load validation set + # reader_bkg = Jets(fn_bkg, nev) # load validation set + # groomed_jets_sig = [] + # for jet in reader_sig.values(): + # groomed_jets_sig.append(dqn.groomer()(jet)) + # masses_sig = np.array(mass(groomed_jets_sig)) + # lower, upper, median = get_window_width(masses_sig) + # groomed_jets_bkg = [] + # for jet in reader_bkg.values(): + # groomed_jets_bkg.append(dqn.groomer()(jet)) + # masses_bkg = np.array(mass(groomed_jets_bkg)) + # # calculate the loss function + # count_bkg = ((masses_bkg > lower) & (masses_bkg < upper)).sum() + # frac_bkg = count_bkg/float(len(masses_bkg)) + # loss = abs(upper-lower)/5 + abs(median-massref) + frac_bkg*20 + # return loss, (lower,upper,median) + +#---------------------------------------------------------------------- +def load_environment(env_setup): + low = np.array( + [ + 0., # squared distance + 0., # E + -0.37260447692861504, # x + -0.35284, # z + 0., # cluster idx + 0., # status + ], dtype=np.float32 + ).reshape(1,-1) + low = np.repeat(low, 1+env_setup['k'], 0) + + high = np.array( + [ + 3., # squared distance + 500., # E + 0.37260447692861504, # x + 0.91702, # z + 1.50, # cluster idx / 100 + 15., # status / 100 + ], dtype=np.float32 + ).reshape(1,-1) + high = np.repeat(high, 1+env_setup['k'], 0) + + slicerl_env = SlicerlEnvContinuous(env_setup, low=low, high=high) + return slicerl_env + +#---------------------------------------------------------------------- +def build_and_train_model(slicerl_agent_setup, slicerl_env=None): + """Run a test model""" + + if slicerl_env is None: + env_setup = slicerl_agent_setup.get('slicerl_env') + slicerl_env = load_environment(env_setup) + + agent_setup = slicerl_agent_setup.get('rl_agent') + ddpg = build_ddpg(agent_setup, slicerl_env.observation_space.shape, slicerl_env.k) + + logdir = '%s/logs/{}'.format(time()) % slicerl_agent_setup['output'] + print(f'[+] Constructing tensorboard log in {logdir}') + tensorboard = TensorBoard(log_dir=logdir) + + print('[+] Fitting DDPG agent...') + r = ddpg.fit(slicerl_env, nb_steps=agent_setup['nstep'], + visualize=False, verbose=1, callbacks=[tensorboard]) + + # compute nominal reward after training + median_reward = np.median(r.history['episode_reward']) + print(f'[+] Median reward: {median_reward}') + + # After training is done, we save the final weights. + if not slicerl_agent_setup['scan']: + weight_file = '%s/weights.h5' % slicerl_agent_setup['output'] + print(f'[+] Saving weights to {weight_file}') + ddpg.save_weights(weight_file, overwrite=True) + + # save the model architecture in json + model_file = '%s/model' % slicerl_agent_setup['output'] + actor_file = '%s_actor.json' % model_file + critic_file = '%s_critic.json' % model_file + print(f'[+] Saving model to {actor_file}, {critic_file}') + with open(actor_file, 'w') as outfile: + json.dump(ddpg.actor.to_json(), outfile) + with open(critic_file, 'w') as outfile: + json.dump(ddpg.critic.to_json(), outfile) + + if slicerl_agent_setup['scan']: + # compute a metric for training set (TODO: change to validation) + raise ValueError('SCAN LOSS FCT NOT IMPLEMENTED YET') + loss, window = loss_calc(ddpg, env_setup['val'], env_setup['nev_val']) + print(f'Loss function for scan = {loss}') + res = {'loss': loss, 'reward': median_reward, 'status': STATUS_OK} + else: + res = ddpg + return res diff --git a/src/slicerl/read_data.py b/src/slicerl/read_data.py new file mode 100644 index 0000000..de49281 --- /dev/null +++ b/src/slicerl/read_data.py @@ -0,0 +1,171 @@ +# This file is part of SliceRL by M. Rossi + +import json, gzip, sys, csv +from abc import ABC, abstractmethod +import numpy as np +from math import log, ceil, floor, pi +from slicerl.Event import Event, CaloHit + +#====================================================================== +class Reader(object): + """ + Reader for files consisting of a sequence of json objects. + Any pure string object is considered to be part of a header (even if it appears at the end!) + """ + + #---------------------------------------------------------------------- + def __init__(self, infile, nmax=-1, num_lines=6): + """Initialize the reader.""" + self.infile = infile + self.readline_fn = lambda x: np.array(next(x))[:-1].astype(np.float32) + self.nmax = nmax + self.num_lines = num_lines + self.reset() + + #---------------------------------------------------------------------- + def reset(self): + """ + Reset the reader to the start of the file, clear the header and event count. + """ + self.stream = csv.reader( gzip.open(self.infile,'rt') ) + self.n = 0 + self.header = [] + + #---------------------------------------------------------------------- + def __iter__(self): + # needed for iteration to work + return self + + #---------------------------------------------------------------------- + def __next__(self): + ev = self.next_event() + if (ev is None): raise StopIteration + else : return ev + + #---------------------------------------------------------------------- + def next(self): return self.__next__() + + #---------------------------------------------------------------------- + def next_event(self): + # we have hit the maximum number of events + if (self.n == self.nmax): + print ("# Exiting after having read nmax events") + return None + + try: + c = [] + c.append( self.readline_fn(self.stream) / 100 ) # energies [ADC]/100 + c.append( self.readline_fn(self.stream) / 1000.) # xs [10^1 m] + c.append( self.readline_fn(self.stream) / 1000.) # zs [10^1 m] + c.append( self.readline_fn(self.stream) ) # cluster_idx + c.append( self.readline_fn(self.stream) ) # pndr_idx + c.append( self.readline_fn(self.stream) ) # cheating_idx (mc truth) + if self.num_lines == 7: + c.append( self.readline_fn(self.stream) ) # slicerl_idx + else: + c.append( -1*np.ones_like(c[-1]) ) + except IOError: + print("# got to end with IOError (maybe gzip structure broken?) around event", self.n, file=sys.stderr) + return None + except EOFError: + print("# got to end with EOFError (maybe gzip structure broken?) around event", self.n, file=sys.stderr) + return None + except ValueError: + print("# got to end with ValueError (empty json entry) around event", self.n, file=sys.stderr) + return None + except StopIteration: + print("# Exiting after having read all the events") + return None + + self.n += 1 + return np.stack(c) + +#====================================================================== +class Image(ABC): + """Image which transforms point-like information into pixelated 2D + images which can be processed by convolutional neural networks.""" + def __init__(self, infile, nmax, num_lines=6): + self.reader = Reader(infile, nmax, num_lines) + + #---------------------------------------------------------------------- + @abstractmethod + def process(self, event): + pass + + #---------------------------------------------------------------------- + def values(self): + res = [] + while True: + event = self.reader.next_event() + if event is not None: + res.append(self.process(event)) + else: + break + self.reader.reset() + return res + +#====================================================================== +class Events(Image): + """Read input file with calohits and transform into python events.""" + + #---------------------------------------------------------------------- + def __init__(self, infile, nmax, k, min_hits, num_lines=None): + """ + Parameters + ---------- + infile : str, input file name + nmax : int, max number of events to load + k : the number of neighbours calohits to include in the Graph + min_hits : int, consider slices with more than min_hits Calohits only + num_lines : int, number of lines stored in the file. Specifies if the + file contains inference results or not. + """ + Image.__init__(self, infile, nmax, num_lines) + self.min_hits = min_hits + self.printouts=10 + self.k = k + + #---------------------------------------------------------------------- + def process(self, event): + return Event(event, self.k, self.min_hits) + +#====================================================================== +def load_Events_from_file(filename, nev, k, min_hits=1, num_lines=6): + """ + Utility function to load Events object from file. Return a list of Event + objects. + + Parameters + ---------- + - filename : str, file to load events from + - nev : int, number of events to load + - k : the number of neighbours calohits to include in the Graph + - min_hits : int, consider slices with more than min_hits Calohits only + - num_lines : int, number of lines stored in the file. Lines stand for: + energies, xs, zs, cluster_idx, pndr_idx, cheating_idx, + slicerl_idx (optional). + + Returns + ------- + - list + list of loaded Event object (with length equal to nev) + """ + reader = Events(filename, nev, k, min_hits, num_lines) + return reader.values() + +#====================================================================== +def save_Event_list_to_file(events, filename): + """ + Utility function to save Event.particles lists to file. Each line represents + an Event, with all Particle objects contained in Event.particles. For each + Particle, store 4-momentum and the additional information contained in PU + and status attributes. + + Parameters + ---------- + - events : list, list of Event objects + - filename : str + """ + with gzip.open(filename, 'wb') as wfp: + for event in events: + event.dump(wfp) diff --git a/src/slicerl/scripts/__init__.py b/src/slicerl/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/slicerl/scripts/slicerl.py b/src/slicerl/scripts/slicerl.py new file mode 100644 index 0000000..7fadb60 --- /dev/null +++ b/src/slicerl/scripts/slicerl.py @@ -0,0 +1,254 @@ +# This file is part of SliceRL by M. Rossi + +""" + slicerl.py: the entry point for slicerl. +""" +from slicerl.read_data import Events, save_Event_list_to_file, load_Events_from_file +from slicerl.Event import Event +from slicerl.models import build_and_train_model, load_runcard, load_environment +from slicerl.diagnostics import inference, make_plots, load_and_dump_plots +from slicerl.Slicer import ContinuousSlicer +from slicerl.keras_to_cpp import keras_to_cpp, check_model +from hyperopt import fmin, tpe, hp, Trials, space_eval +from hyperopt.mongoexp import MongoTrials +from time import time as tm +from shutil import copyfile +from copy import deepcopy +import os, argparse, pickle, pprint, json, gzip, ast, shutil +#import cProfile + +os.environ['CUDA_VISIBLE_DEVICES'] = "" + +#---------------------------------------------------------------------- +def run_hyperparameter_scan(search_space): + """Running a hyperparameter scan using hyperopt. + TODO: implement cross-validation, e.g. k-fold, or randomized cross-validation. + TODO: use test data as hyper. optimization goal. + TODO: better import/export for the best model, wait to DQNAgentSlicerl + """ + + print('[+] Performing hyperparameter scan...') + max_evals = search_space['cluster']['max_evals'] + if search_space['cluster']['enable']: + url = search_space['cluster']['url'] + key = search_space['cluster']['exp_key'] + trials = MongoTrials(url, exp_key=key) + slicerl_env = None + else: + env_setup = search_space.get('slicerl_env') + slicerl_env = load_environment(env_setup) + trials = Trials() + + best = fmin(lambda p: build_and_train_model(p, slicerl_env), search_space, algo=tpe.suggest, max_evals=max_evals, trials=trials) + + best_setup = space_eval(search_space, best) + print('\n[+] Best scan setup:') + pprint.pprint(best_setup) + + log = '%s/hyperopt_log_{}.pickle'.format(tm()) % search_space['output'] + with open(log,'wb') as wfp: + print(f'[+] Saving trials in {log}') + pickle.dump(trials.trials, wfp) + + # disable scan for final fit + best_setup['scan'] = False + return best_setup + +#---------------------------------------------------------------------- +def load_json(runcard_file): + """Loads json, execute python expressions, and sets + scan flags accordingly to the syntax. + """ + runcard = load_runcard(runcard_file) + runcard['scan'] = False + for key, value in runcard.get('slicerl_env').items(): + if 'hp.' in str(value): + runcard['slicerl_env'][key] = eval(value) + runcard['scan'] = True + for key, value in runcard.get('rl_agent').items(): + if 'hp.' in str(value): + runcard['rl_agent'][key] = eval(value) + runcard['scan'] = True + return runcard + +#---------------------------------------------------------------------- +def makedir(folder): + """Create directory.""" + if not os.path.exists(folder): + os.mkdir(folder) + else: + raise Exception('Output folder %s already exists.' % folder) + +#---------------------------------------------------------------------- +def safe_inference_and_plots(slicer, fnin, fnres, plotdir, loaddir, nev, k): + """ + Make inference from dataset at fnin and save results at fnres. Output + diagnostic plots in plotdir and plot data in loaddir. Perform checks trying + to skip already stored computations. This function never overwrites files. + + Parameters + ---------- + fnin: str + input dataset file name + fnres: str + inference results file name + plotdir: str + folder name to store diagnostic plots + loaddir: str + folder name to store diagnostic plots data (to tweak plots only) + nev: int + number of event to load + """ + try: + makedir(plotdir) + makedir(loaddir) + except: + if len(os.listdir(loaddir)) > 0: + print(f'[+] Plotting on previously collected results: {loaddir} already exists') + load_and_dump_plots(plotdir, loaddir) + else: + print(f'[+] Ignoring plots: {loaddir} already exists, but is empty') + else: + print(f'[+] Creating test plots in {plotdir}') + # if already tested slicer, load results and make plots + # otherwise do inference first + if os.path.isfile(fnres): + print('[+] Json file found: loading from saved test data') + events = load_Events_from_file(fnres, nev, k, num_lines=7) + else: + events = load_Events_from_file(fnin, nev, k) + events = inference(slicer, events) + save_Event_list_to_file(events, fnres) + make_plots(events, plotdir) + + +#---------------------------------------------------------------------- +def main(): + """Parsing command line arguments""" + start = tm() + # read command line arguments + parser = argparse.ArgumentParser(description='Train an ML slicer.') + parser.add_argument('runcard', action='store', nargs='?', default=None, + help='A json file with the setup.') + parser.add_argument('--model', '-m', type=str, default=None, + help='The input model.') + parser.add_argument('--output', '-o', type=str, default=None, + help='The output folder.') + parser.add_argument('--plot',action='store_true',dest='plot') + parser.add_argument('--force', '-f', action='store_true',dest='force', + help='Overwrite existing files if present') + parser.add_argument('--cpp',action='store_true',dest='cpp') + parser.add_argument('--data', type=str, default=None, dest='data', + help='Data on which to apply the slicer.') + parser.add_argument('--nev', '-n', type=float, default=-1, + help='Number of events.') + args = parser.parse_args() + + # check that input is coherent + if (not args.model and not args.runcard) or (args.model and args.runcard): + raise ValueError('Invalid options: requires either input runcard or model.') + elif args.runcard and not os.path.isfile(args.runcard): + raise ValueError('Invalid runcard: not a file.') + elif args.model and not (args.plot or args.cpp or args.data): + raise ValueError('Invalid options: no actions requested.') + if args.force: + print('WARNING: Running with --force option will overwrite existing model') + + if args.runcard: + # load json + setup = load_json(args.runcard) + + # create output folder + base = os.path.basename(args.runcard) + out = os.path.splitext(base)[0] + if args.output is not None: + out = args.output + try: + makedir(out) + except Exception as error: + if args.force: + print(f'WARNING: Overwriting {out} with new model') + shutil.rmtree(out) + makedir(out) + else: + print(error) + print('Delete or run with "--force" to overwrite.') + exit(-1) + setup['output'] = out + + # copy runcard to output folder + copyfile(args.runcard, f'{out}/input-runcard.json') + + # slicerl common environment setup + if setup.get('scan'): + rl_agent_setup = run_hyperparameter_scan(setup) + else: + # create the DQN agent and train it. + rl_agent_setup = setup + + print('[+] Training best model:') + ddpg = build_and_train_model(rl_agent_setup) + + # save the final runcard + with open(f'{out}/runcard.json','w') as f: + json.dump(rl_agent_setup, f, indent=4) + + fnres = '%s/test_predictions.json.gz' % setup['output'] + + print('[+] Done with training, now testing on sample set') + if os.path.exists(fnres): + os.remove(fnres) + + slicer = ddpg.slicer() + events = load_Events_from_file(setup['test']['fn'], args.nev, setup['slicerl_env']['k']) + events = inference(slicer, events) + save_Event_list_to_file(events, fnres) + + # define the folder where to do the plotting/cpp conversation + folder = setup['output'] + + elif args.model: + # raise ValueError('Loading of model has not been implemented yet') + folder = args.model.strip('/') + # loading json card + setup = load_runcard('%s/runcard.json' % folder) + rl_agent_setup = setup + # loading slicer + slicer = ContinuousSlicer() + modelwgts_fn = '%s/weights.h5' % folder + modeljson_fn = '%s/model.json' % folder + slicer.load_with_json(modeljson_fn, modelwgts_fn) + + # if requested, add plotting + if args.plot: + fnin = setup['test']['fn'] + plotdir = '%s/plots' % folder + loaddir = '%s/results' % plotdir + fnres = '%s/test_predictions.json.gz' % setup['output'] + safe_inference_and_plots(slicer, fnin, fnres, plotdir, loaddir, args.nev, setup['slicerl_env']['k']) + + # if a data set was given as input, produce plots from it + # always check if inference data is already there + if args.data: + fnin = os.path.basename(args.data).split(os.extsep)[0] + plotdir='%s/%s' % (folder, fnin) + loaddir = '%s/results' % plotdir + fnres = '%s/%s_sliced.json.gz' % (plotdir, fnin) + safe_inference_and_plots(slicer, args.data, fnres, plotdir, loaddir, args.nev, setup['slicerl_env']['k']) + + # if requested, add cpp output + if args.cpp: + check_model(rl_agent_setup['rl_agent']) + cppdir = '%s/cpp' % folder + try: + makedir(cppdir) + except: + print(f'[+] Ignoring cpp instruction: {cppdir} already exists') + else: + print(f'[+] Adding cpp model in {cppdir}') + cpp_fn = '%s/model.nnet' % cppdir + arch_dic=ast.literal_eval(slicer.model.to_json() + .replace('true','True') + .replace('null','None')) + keras_to_cpp(slicer.model, arch_dic['config']['layers'], cpp_fn) + print(f"Program done in {tm()-start} s") diff --git a/src/slicerl/tools.py b/src/slicerl/tools.py new file mode 100644 index 0000000..0314d7a --- /dev/null +++ b/src/slicerl/tools.py @@ -0,0 +1,79 @@ +# This file is part of SliceRL by M. Rossi + +import numpy as np +import math +from energyflow.emd import emd + + +#---------------------------------------------------------------------- +def mass(events, noPU=False): + """Given a list of events, determine the masses for each jet inside.""" + + masses = [] + + for event in events: + for jet_noPU, jet in event.jet_pairs: + # masses_per_event = [] + p = jet_noPU if noPU else jet + msq = p.E**2 - p.px**2 - p.py**2 - p.pz**2 + # masses_per_event.append(math.sqrt(msq) if msq > 0.0 else -math.sqrt(-msq)) + masses.append(math.sqrt(msq) if msq > 0.0 else -math.sqrt(-msq)) + # masses.append(masses_per_event) + + return masses + +#---------------------------------------------------------------------- +def pT(events, noPU=False): + """Given a list of events, determine the pT for each jet inside.""" + + return [math.sqrt(p.px()**2 + p.py()**2) for event in events for jet in event] + +#---------------------------------------------------------------------- +def get_window_width(masses, lower_frac=20, upper_frac=80): + """Returns""" + lower = np.nanpercentile(masses, lower_frac) + upper = np.nanpercentile(masses, upper_frac) + median = np.median(masses[(masses > lower) & (masses < upper)]) + return lower, upper, median + +#---------------------------------------------------------------------- +def quality_metric(slice_state, mc_state): + """ Caluclate EMD metric for two Particle objects """ + return emd(slice_state, mc_state, R=0.6) + +#---------------------------------------------------------------------- +def jet_emd(j_noPU, j): + """ Caluclate EMD metric for two FastJet.PseudoJet objects """ + pT = lambda x: math.sqrt(x.px()**2 + x.py()**2) + ev0 = np.array([pT(j_noPU), j_noPU.rap(), j_noPU.phi()]) + ev1 = np.array([pT(j), j.rap(), j.phi()]) + return emd(ev0, ev1) + +#---------------------------------------------------------------------- +def confusion_matrix_per_event(scores): + """ + Computes confusion matrix values for each jet in list + - scores list of event numpy array scores per event: first level events, + second level arrays of shape [num particles, 2] + Returns + - tp true positives ratio in jet distribution + - fp false positives ratio in jet distribution + - fn false negatives ratio in jet distribution + - tn true negatives ratio in jet distribution + """ + tp = [] + fp = [] + fn = [] + tn = [] + for score in scores: + if score.shape[0] == 0: + continue + truths = score[:,0].astype(bool) + preds = score[:,1] + tot = score.shape[0] + tp.append(len(np.where(preds[truths] == 1)[0])/tot) + fp.append(len(np.where(preds[~truths] == 1)[0])/tot) + fn.append(len(np.where(preds[truths] == 0)[0])/tot) + tn.append(len(np.where(preds[~truths] == 0)[0])/tot) + + return np.array(tp), np.array(fp), np.array(fn), np.array(tn) \ No newline at end of file diff --git a/test_data_03GeV.csv.gz b/test_data_03GeV.csv.gz new file mode 100644 index 0000000..2d2f5f0 Binary files /dev/null and b/test_data_03GeV.csv.gz differ diff --git a/test_data_2GeV.csv.gz b/test_data_2GeV.csv.gz new file mode 100644 index 0000000..7bf74ad Binary files /dev/null and b/test_data_2GeV.csv.gz differ