1283568872

Build Type

github

Committed by GitHub

Commit Message

<a href="https://github.com/XENONnT/straxen/commit/<a class=hub.com/XENONnT/straxen/commit/61bc48165068cf4e925a2e0024ee81d12a86cf7b">61bc48165<a href="https://github.com/XENONnT/straxen/commit/61bc48165068cf4e925a2e0024ee81d12a86cf7b">">Merge </a><a class="double-link" href="https://github.com/XENONnT/straxen/commit/<a class="double-link" href="https://github.com/XENONnT/straxen/commit/8337f5456d82c512b1aded092b15310ac332782b">8337f5456</a>">8337f5456</a><a href="https://github.com/XENONnT/straxen/commit/61bc48165068cf4e925a2e0024ee81d12a86cf7b"> into 0095eb5cc">0095eb5cc</a>

Pull Request Pull Request #671: S1 posrec events

Run Details

95 of 95 new or added lines in 1 file covered. (100.0%)

5192 of 6618 relevant lines covered (78.45%)

0.78 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/straxen/plugins/s1_position_reconstruction.py

"""S1-based 3D position reconstruction for Xenon-nT"""

import os
import tempfile
import tarfile
import numpy as np
import strax
import straxen
from warnings import warn
export, __all__ = strax.exporter()

# Current defaul algorithm: data-driven CNN. 
DEFAULT_S1POSREC_ALGO_OPTION = tuple([strax.Option("default_s1reconstruction_algorithm",
                 help="default S1 reconstruction algorithm that provides (x,y,z)",
                 default="cnn_s1",
                 )])

@export
@strax.takes_config(
        strax.Option("s1_cnn_model_path", 
                     help="Path to the S1 CNN model", 
                     default=("/project2/lgrandi/guidam/CNN_S1_XYZ_SAVED_MODELS/version_datadriven_00_080921")
                    ),
    strax.Option('n_top_pmts', default=straxen.n_top_pmts,
                 help="Number of top PMTs")
)


class S1EventPositionBase(strax.Plugin):
    """
    Base class for S1 three-dimensional position reconstruction.
    This class should only be used when subclassed for the different algorithms. 
    """

    depends_on = ('events','event_area_per_channel',)
    algorithm = None
    compressor = 'zstd'
    parallel = True
    __version__ = '0.0.0'

    def infer_dtype(self):
        if self.algorithm is None:
            raise NotImplementedError(f'Base class should not be used without '
                                      f'algorithm as done in {__class__.__name__}')
        dtype = [('x_' + self.algorithm, np.float32,
                  f'Reconstructed {self.algorithm} S1 X position (cm), uncorrected'),
                 ('y_' + self.algorithm, np.float32,
                  f'Reconstructed {self.algorithm} S1 Y position (cm), uncorrected'),
                 ('z_' + self.algorithm, np.float32,
                  f'Reconstructed {self.algorithm} S1 Z position (cm), uncorrected')
                  ]

        dtype += strax.time_fields
        return dtype

    def setup(self):
        self.model_file = self._get_model_file_name()
        if self.model_file is None:
            warn(f'No file provided for {self.algorithm}. Setting all values '
                 f'for {self.provides} to None.')
            # No further setup required
            return

        # Load the tensorflow model
        import tensorflow as tf
        if os.path.exists(self.model_file):
            print(f"Path is local. Loading {self.algorithm} TF model locally "
                  f"from disk.")
        else:
            downloader = straxen.MongoDownloader()
            try:
                self.model_file = downloader.download_single(self.model_file)
            except straxen.mongo_storage.CouldNotLoadError as e:
                raise RuntimeError(f'Model files {self.model_file} is not found') from e
        with tempfile.TemporaryDirectory() as tmpdirname:
        # Allow to load the model both from folder or from file.
            try:
                tar = tarfile.open(self.model_file, mode="r:gz")
                tar.extractall(path=tmpdirname)
                self.model = tf.keras.models.load_model(tmpdirname)
            except:
                self.model = tf.keras.models.load_model(self.model_file)


    def compute(self,events):
        result = np.ones(len(events), dtype=self.dtype)
        result['time'], result['endtime'] = events['time'], strax.endtime(events)

        result['x_' + self.algorithm] *= float('nan')
        result['y_' + self.algorithm] *= float('nan')
        result['z_' + self.algorithm] *= float('nan')

        if self.model_file is None:
            # This plugin is disabled since no model is provided
            return result


        # Getting actual position reconstruction
        _in = events['s1_area_per_channel']
        with np.errstate(divide='ignore', invalid='ignore'):
        # Normalise patters by dividing by largest PMT output between the two arrays. 
            _in = _in / _in.max(axis=1,keepdims=True)
        _out = self.model.predict(_in)

        # writing output to the result
        
        result['x_' + self.algorithm] = _out[:, 0]
        result['y_' + self.algorithm] = _out[:, 1]
        result['z_' + self.algorithm] = _out[:, 2]
        return result

    def _get_model_file_name(self):

        config_file = f'{self.algorithm}_model'
        model_from_config = self.config.get(config_file, 'No file')
        if model_from_config == 'No file':
            raise ValueError(f'{__class__.__name__} should have {config_file} '
                             f'provided as an option.')
        if isinstance(model_from_config, str) and os.path.exists(model_from_config):
            # Allow direct path specification
            return model_from_config
        if model_from_config is None:
            # Allow None to be specified (disables processing for given posrec)
            return model_from_config

        # Use CMT
        model_file = straxen.get_config_from_cmt(self.run_id, model_from_config)
        return model_file

@export
@strax.takes_config(
    strax.Option('cnn_s1_model',
                 help='Neural network model.' 
                      'If CMT, specify as (CMT_model, (cnn_s1_model, ONLINE), True)))'
                      'Set to None to skip the computation of this plugin.',
                default="/project2/lgrandi/guidam/CNN_S1_XYZ_SAVED_MODELS/version00_020521"
#                  default=("CMT_model", ('cnn_s1_model', "None"), True)
                )
)

class S1EventPositionCNN(S1EventPositionBase):
    """Convolutional Neural Network (CNN) neural net for S1 (x,y,z) position reconstruction"""
    provides = "s1_event_positions_cnn"
    algorithm = "cnn_s1"
    __version__ = '0.0.1'


@export
@strax.takes_config(
    *DEFAULT_S1POSREC_ALGO_OPTION
)
class S1EventPosition(strax.MergeOnlyPlugin):
    """
    Merge the reconstructed S1 algorithms of the different algorithms.
    
    Select one of the plugins to provide the 'x','y' and 'z' to be used 
    further down the chain. Since we already have the information
    needed here, there is no need to wait until events to make the
    decision.
        
    Since the computation is trivial as it only combined the three 
    input plugins, don't save this plugins output.
    """
    provides = "s1_event_positions"
    depends_on = ("s1_event_positions_cnn")
    save_when = strax.SaveWhen.NEVER
    __version__ = '0.0.0'

    def infer_dtype(self):
        dtype = strax.merged_dtype([self.deps[d].dtype_for(d) for d in self.depends_on])
        dtype += [('x', np.float32, 'Reconstructed S1 X position (cm), uncorrected'),
                  ('y', np.float32, 'Reconstructed S1 Y position (cm), uncorrected'),
                  ('z', np.float32, 'Reconstructed S1 Z position (cm), uncorrected')]
        return dtype

    def compute(self, events):
        result = {dtype: events[dtype] for dtype in events.dtype.names}
        algorithm = self.config['default_s1reconstruction_algorithm']
        if not 'x_' + algorithm in events.dtype.names:
            raise ValueError
        for xyz in ('x', 'y', 'z'):
            result[xyz] = events[f'{xyz}_{algorithm}']
        return result


1	"""S1-based 3D position reconstruction for Xenon-nT"""
2
3	import os	×
4	import tempfile	×
5	import tarfile	×
6	import numpy as np	×
7	import strax	×
8	import straxen	×
9	from warnings import warn	×
10	export, __all__ = strax.exporter()	×
11
12	# Current defaul algorithm: data-driven CNN.
13	DEFAULT_S1POSREC_ALGO_OPTION = tuple([strax.Option("default_s1reconstruction_algorithm",	×
14	help="default S1 reconstruction algorithm that provides (x,y,z)",
15	default="cnn_s1",
16	)])
17
18	@export	×
19	@strax.takes_config(	×
20	strax.Option("s1_cnn_model_path",
21	help="Path to the S1 CNN model",
22	default=("/project2/lgrandi/guidam/CNN_S1_XYZ_SAVED_MODELS/version_datadriven_00_080921")
23	),
24	strax.Option('n_top_pmts', default=straxen.n_top_pmts,
25	help="Number of top PMTs")
26	)
27
28
29	class S1EventPositionBase(strax.Plugin):	×
30	"""
31	Base class for S1 three-dimensional position reconstruction.
32	This class should only be used when subclassed for the different algorithms.
33	"""
34
35	depends_on = ('events','event_area_per_channel',)	×
36	algorithm = None	×
37	compressor = 'zstd'	×
38	parallel = True	×
39	__version__ = '0.0.0'	×
40
41	def infer_dtype(self):	×
42	if self.algorithm is None:	×
43	raise NotImplementedError(f'Base class should not be used without '	×
44	f'algorithm as done in {__class__.__name__}')
45	dtype = [('x_' + self.algorithm, np.float32,	×
46	f'Reconstructed {self.algorithm} S1 X position (cm), uncorrected'),
47	('y_' + self.algorithm, np.float32,
48	f'Reconstructed {self.algorithm} S1 Y position (cm), uncorrected'),
49	('z_' + self.algorithm, np.float32,
50	f'Reconstructed {self.algorithm} S1 Z position (cm), uncorrected')
51	]
52
53	dtype += strax.time_fields	×
54	return dtype	×
55
56	def setup(self):	×
57	self.model_file = self._get_model_file_name()	×
58	if self.model_file is None:	×
59	warn(f'No file provided for {self.algorithm}. Setting all values '	×
60	f'for {self.provides} to None.')
61	# No further setup required
62	return	×
63
64	# Load the tensorflow model
65	import tensorflow as tf	×
66	if os.path.exists(self.model_file):	×
67	print(f"Path is local. Loading {self.algorithm} TF model locally "	×
68	f"from disk.")
69	else:
70	downloader = straxen.MongoDownloader()	×
71	try:	×
72	self.model_file = downloader.download_single(self.model_file)	×
73	except straxen.mongo_storage.CouldNotLoadError as e:	×
74	raise RuntimeError(f'Model files {self.model_file} is not found') from e	×
75	with tempfile.TemporaryDirectory() as tmpdirname:	×
76	# Allow to load the model both from folder or from file.
77	try:	×
78	tar = tarfile.open(self.model_file, mode="r:gz")	×
79	tar.extractall(path=tmpdirname)	×
80	self.model = tf.keras.models.load_model(tmpdirname)	×
81	except:	×
82	self.model = tf.keras.models.load_model(self.model_file)	×
83
84
85	def compute(self,events):	×
86	result = np.ones(len(events), dtype=self.dtype)	×
87	result['time'], result['endtime'] = events['time'], strax.endtime(events)	×
88
89	result['x_' + self.algorithm] *= float('nan')	×
90	result['y_' + self.algorithm] *= float('nan')	×
91	result['z_' + self.algorithm] *= float('nan')	×
92
93	if self.model_file is None:	×
94	# This plugin is disabled since no model is provided
95	return result	×
96
97
98	# Getting actual position reconstruction
99	_in = events['s1_area_per_channel']	×
100	with np.errstate(divide='ignore', invalid='ignore'):	×
101	# Normalise patters by dividing by largest PMT output between the two arrays.
102	_in = _in / _in.max(axis=1,keepdims=True)	×
103	_out = self.model.predict(_in)	×
104
105	# writing output to the result
106
107	result['x_' + self.algorithm] = _out[:, 0]	×
108	result['y_' + self.algorithm] = _out[:, 1]	×
109	result['z_' + self.algorithm] = _out[:, 2]	×
110	return result	×
111
112	def _get_model_file_name(self):	×
113
114	config_file = f'{self.algorithm}_model'	×
115	model_from_config = self.config.get(config_file, 'No file')	×
116	if model_from_config == 'No file':	×
117	raise ValueError(f'{__class__.__name__} should have {config_file} '	×
118	f'provided as an option.')
119	if isinstance(model_from_config, str) and os.path.exists(model_from_config):	×
120	# Allow direct path specification
121	return model_from_config	×
122	if model_from_config is None:	×
123	# Allow None to be specified (disables processing for given posrec)
124	return model_from_config	×
125
126	# Use CMT
127	model_file = straxen.get_config_from_cmt(self.run_id, model_from_config)	×
128	return model_file	×
129
130	@export	×
131	@strax.takes_config(	×
132	strax.Option('cnn_s1_model',
133	help='Neural network model.'
134	'If CMT, specify as (CMT_model, (cnn_s1_model, ONLINE), True)))'
135	'Set to None to skip the computation of this plugin.',
136	default="/project2/lgrandi/guidam/CNN_S1_XYZ_SAVED_MODELS/version00_020521"
137	# default=("CMT_model", ('cnn_s1_model', "None"), True)
138	)
139	)
140
141	class S1EventPositionCNN(S1EventPositionBase):	×
142	"""Convolutional Neural Network (CNN) neural net for S1 (x,y,z) position reconstruction"""
143	provides = "s1_event_positions_cnn"	×
144	algorithm = "cnn_s1"	×
145	__version__ = '0.0.1'	×
146
147
148	@export	×
149	@strax.takes_config(	×
150	*DEFAULT_S1POSREC_ALGO_OPTION
151	)
152	class S1EventPosition(strax.MergeOnlyPlugin):	×
153	"""
154	Merge the reconstructed S1 algorithms of the different algorithms.
155
156	Select one of the plugins to provide the 'x','y' and 'z' to be used
157	further down the chain. Since we already have the information
158	needed here, there is no need to wait until events to make the
159	decision.
160
161	Since the computation is trivial as it only combined the three
162	input plugins, don't save this plugins output.
163	"""
164	provides = "s1_event_positions"	×
165	depends_on = ("s1_event_positions_cnn")	×
166	save_when = strax.SaveWhen.NEVER	×
167	__version__ = '0.0.0'	×
168
169	def infer_dtype(self):	×
170	dtype = strax.merged_dtype([self.deps[d].dtype_for(d) for d in self.depends_on])	×
171	dtype += [('x', np.float32, 'Reconstructed S1 X position (cm), uncorrected'),	×
172	('y', np.float32, 'Reconstructed S1 Y position (cm), uncorrected'),
173	('z', np.float32, 'Reconstructed S1 Z position (cm), uncorrected')]
174	return dtype	×
175
176	def compute(self, events):	×
177	result = {dtype: events[dtype] for dtype in events.dtype.names}	×
178	algorithm = self.config['default_s1reconstruction_algorithm']	×
179	if not 'x_' + algorithm in events.dtype.names:	×
180	raise ValueError	×
181	for xyz in ('x', 'y', 'z'):	×
182	result[xyz] = events[f'{xyz}_{algorithm}']	×
183	return result	×
184

XENONnT / straxen / 1283568872

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous