def __init__(
self,
sn_model: SNModel,
catalog: VariableCatalog = None,
num_processes: int = 1,
add_scatter: bool = True,
fixed_snr: Optional[float] = None,
abs_mb: float = const.betoule_abs_mb,
cosmo: Cosmology = const.betoule_cosmo
) -> None:
"""Fit light-curves using multiple processes and combine results into an output file
Args:
sn_model: Model to use when simulating light-curves
catalog: Optional reference start catalog to calibrate simulated flux values to
num_processes: Number of processes to allocate to the node
abs_mb: The absolute B-band magnitude of the simulated SNe
cosmo: Cosmology to assume in the simulation
"""
self.sim_model = copy(sn_model)
self.catalog = catalog
self.add_scatter = add_scatter
self.fixed_snr = fixed_snr
self.abs_mb = abs_mb
self.cosmo = cosmo
# Node connectors
self.input = Input('Simulated Cadence')
self.success_output = Output('Simulation Success')
self.failure_output = Output('Simulation Failure')
super().__init__(num_processes=num_processes)
def test_error_if_malformed(self) -> None:
"""Test a malformed error is raised by the ``validate`` method"""
self.test_class.input = Input()
self.test_class.output = Output()
with self.assertRaises(self.malformed_exception):
self.test_class.validate()
def test_overwrite_error_on_connection_overwrite(self) -> None:
"""Test an error is raised when trying to overwrite an existing connection"""
input = Input()
self.output_connector.connect(input)
with self.assertRaises(exceptions.OverwriteConnectionError):
self.output_connector.connect(input)
def setUp(self) -> None:
"""Connect two outputs to a single input"""
self.input = Input()
self.output1 = Output()
self.output2 = Output()
self.output1.connect(self.input)
self.output2.connect(self.input)
def test_stores_value_in_queue(self) -> None:
"""Test the ``put`` method puts data into the underlying connected queue"""
# Create a node with an output connector
output = Output()
input = Input()
output.connect(input)
test_val = 'test_val'
output.put(test_val)
self.assertEqual(input._queue.get(), test_val)
def __init__(self, out_path: Union[str, Path], num_processes=1) -> None:
"""Output node for writing HDF5 data to disk
This node can only be run using a single process. This can be the main
process (``num_processes=0``) or a single forked process (``num_processes=1``.)
Args:
out_path: Path to write data to in HDF5 format
"""
if num_processes not in (0, 1):
raise RuntimeError('Number of processes for ``LoadPlasticcCadence`` must be 0 or 1.')
# Make true to raise errors instead of converting them to warnings
self.debug = False
self.out_path = Path(out_path)
self.input = Input('Data To Write')
self.file_store: Optional[pd.HDFStore] = None
super().__init__(num_processes=num_processes)
class MockNode(MockObject, nodes.Node):
"""A ``Node`` subclass that implements placeholder functions for abstract methods"""
def __init__(self, name: Optional[str] = None) -> None:
self.output = Output()
self.input = Input()
super().__init__(name)
def action(self) -> None: # pragma: no cover
"""Placeholder function to satisfy requirements of abstract parent"""
for x in self.input.iter_get():
self.output.put(x)
class MockTarget(MockObject, nodes.Target):
"""A ``Target`` subclass that implements placeholder functions for abstract methods"""
def __init__(self, name: Optional[str] = None) -> None:
self.input = Input()
self.accumulated_data = []
super().__init__(name)
def action(self) -> None:
"""Placeholder function to satisfy requirements of abstract parent"""
for x in self.input.iter_get():
self.accumulated_data.append(x)
def __init__(
self, sn_model: SNModel, vparams: List[str], bounds: Dict = None, num_processes: int = 1
) -> None:
"""Fit light-curves using multiple processes and combine results into an output file
Args:
sn_model: Model to use when fitting light-curves
vparams: List of parameter names to vary in the fit
bounds: Bounds to impose on ``fit_model`` parameters when fitting light-curves
num_processes: Number of processes to allocate to the node
"""
self.sn_model = sn_model
self.vparams = vparams
self.bounds = bounds
# Node Connectors
self.input = Input('Simulated Light-Curve')
self.success_output = Output('Fitting Success')
self.failure_output = Output('Fitting Failure')
super(FitLightCurves, self).__init__(num_processes=num_processes)
class QueueProperties(TestCase):
"""Test test the exposure of queue properties by the overlying ``Connector`` class"""
def setUp(self) -> None:
self.connector = Input(maxsize=1)
def test_size_matches_queue_size(self) -> None:
"""Test the ``size`` method returns the size of the queue`"""
self.assertEqual(self.connector.size(), 0)
self.connector._queue.put(1)
self.assertEqual(self.connector.size(), 1)
def test_full_state(self) -> None:
"""Test the ``full`` method returns the state of the queue"""
self.assertFalse(self.connector.is_full())
self.connector._queue.put(1)
self.assertTrue(self.connector.is_full())
def test_empty_state(self) -> None:
"""Test the ``empty`` method returns the state of the queue"""
self.assertTrue(self.connector.is_empty())
self.connector._queue.put(1)
# The value of Queue.is_empty() updates asynchronously
time.sleep(1)
self.assertFalse(self.connector.is_empty())
class InstanceDisconnect(TestCase):
"""Test the disconnection of two connectors"""
def setUp(self) -> None:
self.input = Input()
self.output = Output()
self.output.connect(self.input)
def test_both_connectors_are_disconnected(self) -> None:
"""Test both connectors are no longer listed as partners"""
self.assertTrue(self.input.is_connected())
self.assertTrue(self.output.is_connected())
self.output.disconnect(self.input)
self.assertNotIn(self.output, self.input.partners)
self.assertFalse(self.input.is_connected())
self.assertFalse(self.output.is_connected())
def test_error_if_not_connected(self) -> None:
"""Test an error is raised when disconnecting a connector that is not connected"""
with self.assertRaises(MissingConnectionError):
Output().disconnect(Input())
class InputGet(TestCase):
"""Test data retrieval from ``Input`` instances"""
def setUp(self) -> None:
self.input_connector = Input()
def test_error_on_non_positive_refresh(self) -> None:
"""Test a ValueError is raised when ``refresh_interval`` is not a positive number"""
with self.assertRaises(ValueError):
self.input_connector.get(timeout=15, refresh_interval=0)
with self.assertRaises(ValueError):
self.input_connector.get(timeout=15, refresh_interval=-1)
def test_returns_queue_value(self) -> None:
"""Test the ``get`` method retrieves data from the underlying queue"""
test_val = 'test_val'
self.input_connector._queue.put(test_val)
time.sleep(1)
self.assertEqual(test_val, self.input_connector.get(timeout=1000))
def test_timeout_raises_timeout_error(self) -> None:
"""Test a ``TimeoutError`` is raise on timeout"""
with self.assertRaises(TimeoutError):
self.input_connector.get(timeout=1)
def test_kill_signal_on_finished_parent_node(self) -> None:
"""Test a kill signal is returned if the parent node is finished"""
target = MockTarget()
self.assertFalse(target.is_expecting_data())
self.assertIs(target.input.get(timeout=15), KillSignal)
class WritePipelinePacket(Target):
"""Pipeline node for writing pipeline packets to disk
Connectors:
input: A pipeline packet
"""
def __init__(self, out_path: Union[str, Path], num_processes=1) -> None:
"""Output node for writing HDF5 data to disk
This node can only be run using a single process. This can be the main
process (``num_processes=0``) or a single forked process (``num_processes=1``.)
Args:
out_path: Path to write data to in HDF5 format
"""
if num_processes not in (0, 1):
raise RuntimeError('Number of processes for ``LoadPlasticcCadence`` must be 0 or 1.')
# Make true to raise errors instead of converting them to warnings
self.debug = False
self.out_path = Path(out_path)
self.input = Input('Data To Write')
self.file_store: Optional[pd.HDFStore] = None
super().__init__(num_processes=num_processes)
def write_packet(self, packet: PipelinePacket) -> None:
"""Write a pipeline packet to the output file"""
# We are taking the simulated parameters as guaranteed to exist
self.file_store.append('simulation/params', packet.sim_params_to_pandas())
self.file_store.append('message', packet.packet_status_to_pandas().astype(str), min_itemsize={'message': 250})
if packet.light_curve is not None: # else: simulation failed
self.file_store.put(f'simulation/lcs/{packet.snid}', packet.light_curve.to_pandas())
if packet.fit_result is not None: # else: fit failed
self.file_store.append('fitting/params', packet.fitted_params_to_pandas())
if packet.covariance is not None:
self.file_store.put(f'fitting/covariance/{packet.snid}', packet.covariance)
def setup(self) -> None:
"""Open a file accessor object"""
self.file_store = pd.HDFStore(self.out_path, mode='w')
def teardown(self) -> None:
"""Close any open file accessors"""
self.file_store.close()
self.file_store = None
def action(self) -> None:
"""Write data from the input connector to disk"""
for packet in self.input.iter_get():
try:
self.write_packet(packet)
except Exception as excep:
if self.debug:
raise
warnings.warn(f'{self.__class__.__name__}: {repr(excep)}')
class SimulateLightCurves(Node):
"""Pipeline node for simulating light-curves based on PLAsTICC cadences
Connectors:
input: A Pipeline Packet
success_output: Emits pipeline packets successfully decorated with a simulated light-curve
failure_output: Emits pipeline packets for cases where the simulation procedure failed
"""
def __init__(
self,
sn_model: SNModel,
catalog: VariableCatalog = None,
num_processes: int = 1,
add_scatter: bool = True,
fixed_snr: Optional[float] = None,
abs_mb: float = const.betoule_abs_mb,
cosmo: Cosmology = const.betoule_cosmo
) -> None:
"""Fit light-curves using multiple processes and combine results into an output file
Args:
sn_model: Model to use when simulating light-curves
catalog: Optional reference start catalog to calibrate simulated flux values to
num_processes: Number of processes to allocate to the node
abs_mb: The absolute B-band magnitude of the simulated SNe
cosmo: Cosmology to assume in the simulation
"""
self.sim_model = copy(sn_model)
self.catalog = catalog
self.add_scatter = add_scatter
self.fixed_snr = fixed_snr
self.abs_mb = abs_mb
self.cosmo = cosmo
# Node connectors
self.input = Input('Simulated Cadence')
self.success_output = Output('Simulation Success')
self.failure_output = Output('Simulation Failure')
super().__init__(num_processes=num_processes)
def simulate_lc(self, params: Dict[str, float], cadence: ObservedCadence) -> LightCurve:
"""Duplicate a plastic light-curve using the simulation model
Args:
params: The simulation parameters to use with ``self.model``
cadence: The observed cadence of the returned light-curve
"""
# Set model parameters and scale the source brightness to the desired intrinsic brightness
model_for_sim = copy(self.sim_model)
model_for_sim.update({p: v for p, v in params.items() if p in model_for_sim.param_names})
model_for_sim.set_source_peakabsmag(self.abs_mb, 'standard::b', 'AB', cosmo=self.cosmo)
# Simulate the light-curve. Make sure to include model parameters as meta data
duplicated = model_for_sim.simulate_lc(cadence, scatter=self.add_scatter, fixed_snr=self.fixed_snr)
# Rescale the light-curve using the reference star catalog if provided
if self.catalog is not None:
duplicated = self.catalog.calibrate_lc(duplicated, ra=params['ra'], dec=params['dec'])
return duplicated
def action(self) -> None:
"""Simulate light-curves with atmospheric effects"""
for packet in self.input.iter_get():
try:
packet.light_curve = self.simulate_lc(
packet.sim_params, packet.cadence
)
except Exception as excep:
packet.message = f'{self.__class__.__name__}: {repr(excep)}'
self.failure_output.put(packet)
else:
self.success_output.put(packet)
def setUp(self) -> None:
self.input_connector = Input()
self.output_connector = Output()
def __init__(self, name: Optional[str] = None) -> None:
self.input = Input()
self.accumulated_data = []
super().__init__(name)
def __init__(self, name: Optional[str] = None) -> None:
self.output = Output()
self.input = Input()
super().__init__(name)
def setUp(self) -> None:
self.connector = Input(maxsize=10)
def test_raises_missing_connection_with_no_parent(self) -> None:
"""Test the iterator exits if input has no paren"""
with self.assertRaises(MissingConnectionError):
next(Input().iter_get())
def setUp(self) -> None:
self.input = Input()
self.output = Output()
self.output.connect(self.input)
def test_error_if_not_connected(self) -> None:
"""Test an error is raised when disconnecting a connector that is not connected"""
with self.assertRaises(MissingConnectionError):
Output().disconnect(Input())
class FitLightCurves(Node):
"""Pipeline node for fitting simulated light-curves
Connectors:
input: A Pipeline Packet
success_output: Emits pipeline packets with successful fit results
failure_output: Emits pipeline packets for cases where the fitting procedure failed
"""
def __init__(
self, sn_model: SNModel, vparams: List[str], bounds: Dict = None, num_processes: int = 1
) -> None:
"""Fit light-curves using multiple processes and combine results into an output file
Args:
sn_model: Model to use when fitting light-curves
vparams: List of parameter names to vary in the fit
bounds: Bounds to impose on ``fit_model`` parameters when fitting light-curves
num_processes: Number of processes to allocate to the node
"""
self.sn_model = sn_model
self.vparams = vparams
self.bounds = bounds
# Node Connectors
self.input = Input('Simulated Light-Curve')
self.success_output = Output('Fitting Success')
self.failure_output = Output('Fitting Failure')
super(FitLightCurves, self).__init__(num_processes=num_processes)
def fit_lc(self, light_curve: LightCurve, initial_guess: Dict[str, float]) -> Tuple[SNFitResult, SNModel]:
"""Fit the given light-curve
Args:
light_curve: The light-curve to fit
initial_guess: Parameters to use as the initial guess in the chi-squared minimization
Returns:
- The optimization result
- A copy of the model with parameter values set to minimize the chi-square
"""
# Use the true light-curve parameters as the initial guess
model = copy(self.sn_model)
model.update({k: v for k, v in initial_guess.items() if k in self.sn_model.param_names})
return model.fit_lc(
light_curve, self.vparams, bounds=self.bounds,
guess_t0=False, guess_amplitude=False, guess_z=False)
def action(self) -> None:
"""Fit light-curves"""
for packet in self.input.iter_get():
try:
packet.fit_result, packet.fitted_model = self.fit_lc(packet.light_curve, packet.sim_params)
packet.covariance = packet.fit_result.salt_covariance_linear()
except Exception as excep:
packet.message = f'{self.__class__.__name__}: {repr(excep)}'
self.failure_output.put(packet)
else:
packet.message = f'{self.__class__.__name__}: {packet.fit_result.message}'
self.success_output.put(packet)