sgnl.sinks.stillsuit_sink

A sink element to write triggers into a sqlite database.

StillSuitSink dataclass

Bases: SnapShotControlSinkElement, ParallelizeSinkElement

Write triggers into a sqlite database

Collects the triggers arriving on the trigger pad, together with the segments arriving on the segment pads, and writes them out per subbank through StillSuit. Offline the database is written at EOS; online it is snapshotted as the analysis runs.

Parameters:

Name	Type	Description	Default
ifos	list	list, the interferometers taking part in the analysis	None
config_name	str	str, the StillSuit configuration naming the database schema. Required	None
bankids_map	dict[str, list]	dict[str, list], the subbank ids belonging to each bank	None
template_ids	Sequence[Any]	Sequence[Any], the template ids in the bank	None
template_sngls	list	list, per-subbank mappings of template id to the sngl_inspiral row for that template	None
subbankids	Sequence[Any]	Sequence[Any], the subbanks this element writes for	None
itacacac_pad_name	str	str, the name of the sink pad carrying triggers	None
segments_pad_map	dict[str, str]	dict[str, str], the sink pad carrying segments for each interferometer, keyed by interferometer	None
trigger_output	dict[str, str]	dict[str, str], the output database path per subbank. Required offline, where there is no snapshotting	None
process_params	dict | None	dict | None, the parameters recorded in the process table	None
program	str	str, the program name recorded in the process table	''
injection_list	list | None	list | None, the injection files recorded in the simulation table	None
is_online	bool	bool, whether to snapshot the database as the analysis runs rather than writing it at EOS	False
multiprocess	bool	bool, whether this element runs in its own process rather than a thread	False
jobid	int	int, the job identifier recorded alongside the output	0
nsubbank_pretend	bool	bool, whether to write every subbank from a single job	False
verbose	bool	bool, whether to report progress as triggers are written	False
injections	bool	bool, whether this is an injection job	False
strike_object	StrikeObject | None	StrikeObject | None, the ranking statistics to record alongside the triggers, where available	None

Source code in sgnl/sinks/stillsuit_sink.py

@dataclass
class StillSuitSink(SnapShotControlSinkElement, ParallelizeSinkElement):
    """Write triggers into a sqlite database

    Collects the triggers arriving on the trigger pad, together with the
    segments arriving on the segment pads, and writes them out per
    subbank through StillSuit. Offline the database is written at EOS;
    online it is snapshotted as the analysis runs.

    Args:
        ifos:
            list, the interferometers taking part in the analysis
        config_name:
            str, the StillSuit configuration naming the database schema.
            Required
        bankids_map:
            dict[str, list], the subbank ids belonging to each bank
        template_ids:
            Sequence[Any], the template ids in the bank
        template_sngls:
            list, per-subbank mappings of template id to the
            sngl_inspiral row for that template
        subbankids:
            Sequence[Any], the subbanks this element writes for
        itacacac_pad_name:
            str, the name of the sink pad carrying triggers
        segments_pad_map:
            dict[str, str], the sink pad carrying segments for each
            interferometer, keyed by interferometer
        trigger_output:
            dict[str, str], the output database path per subbank.
            Required offline, where there is no snapshotting
        process_params:
            dict | None, the parameters recorded in the process table
        program:
            str, the program name recorded in the process table
        injection_list:
            list | None, the injection files recorded in the simulation
            table
        is_online:
            bool, whether to snapshot the database as the analysis runs
            rather than writing it at EOS
        multiprocess:
            bool, whether this element runs in its own process rather
            than a thread
        jobid:
            int, the job identifier recorded alongside the output
        nsubbank_pretend:
            bool, whether to write every subbank from a single job
        verbose:
            bool, whether to report progress as triggers are written
        injections:
            bool, whether this is an injection job
        strike_object:
            StrikeObject | None, the ranking statistics to record
            alongside the triggers, where available
    """

    ifos: list = None  # type: ignore[assignment]
    config_name: str = None  # type: ignore[assignment]
    bankids_map: dict[str, list] = None  # type: ignore[assignment]
    template_ids: Sequence[Any] = None  # type: ignore[assignment]
    template_sngls: list = None  # type: ignore[assignment]
    subbankids: Sequence[Any] = None  # type: ignore[assignment]
    itacacac_pad_name: str = None  # type: ignore[assignment]
    segments_pad_map: dict[str, str] = None  # type: ignore[assignment]
    trigger_output: dict[str, str] = None  # type: ignore[assignment]
    process_params: dict | None = None  # type: ignore[assignment]
    program: str = ""
    injection_list: list | None = None  # type: ignore[assignment]
    is_online: bool = False
    multiprocess: bool = False
    jobid: int = 0
    nsubbank_pretend: bool = False
    verbose: bool = False
    injections: bool = False
    strike_object: StrikeObject | None = None

    def __post_init__(self):
        # Initialize mutable defaults
        if self.trigger_output is None:
            self.trigger_output = {}
        if self.process_params is None:
            self.process_params = {}
        if self.injection_list is None:
            self.injection_list = []

        self._use_threading_override = not self.multiprocess
        if self.config_name is None:
            raise ValueError("Must provide config_name")
        if not self.is_online and not self.trigger_output:
            raise ValueError("Must provide trigger_output")

        SnapShotControlSinkElement.__post_init__(self)

        self.tables = ["trigger", "event"]
        self.event_dict = {t: [] for t in self.tables}
        self.bankids = self.bankids_map.keys()

        with open(self.config_name) as f:
            self.config = yaml.safe_load(f)

        for bankid in self.bankids:
            # FIXME: use job_tag?? but what to use with multi-bank mode?
            if self.injections:
                fn_bankid = bankid + "_inj"
            else:
                fn_bankid = bankid + "_noninj"

            self.add_snapshot_filename("%s_SGNL_TRIGGERS" % fn_bankid, "sqlite.gz")
        self.register_snapshot()

        #
        # Process
        #
        self.process_row = init_config_row(self.config["process"])
        self.process_row["ifos"] = ",".join(self.ifos)
        self.process_row["is_online"] = int(self.is_online)
        self.process_row["node"] = socket.gethostname()
        self.process_row["program"] = self.program
        self.process_row["unix_procid"] = os.getpid()
        self.process_row["username"] = self.get_username()

        #
        # Process params
        #
        self.params = []
        if self.process_params is not None:
            for name, values in self.process_params.items():
                name = "--%s" % name.replace("_", "-")
                if values is None:
                    continue
                elif isinstance(values, list):
                    for v in values:
                        param_row = init_config_row(self.config["process_params"])
                        param_row["param"] = name
                        param_row["program"] = self.program
                        param_row["value"] = str(v)
                        self.params.append(param_row)
                    continue

                param_row = init_config_row(self.config["process_params"])
                param_row["param"] = name
                param_row["program"] = self.program
                param_row["value"] = str(values)
                self.params.append(param_row)

        #
        # Filter
        #
        self.filters = {k: [] for k in self.bankids}
        if self.nsubbank_pretend:
            subbank = self.template_sngls[0]
            for template_id, sngl in subbank.items():
                filter_row = init_config_row(self.config["filter"])
                filter_row["_filter_id"] = template_id
                filter_row["bank_id"] = int(self.subbankids[0].split("_")[0])
                filter_row["subbank_id"] = int(self.subbankids[0].split("_")[1])
                filter_row["end_time_delta"] = (
                    sngl.end_time * 1_000_000_000 + sngl.end_time_ns
                )
                filter_row["mass1"] = sngl.mass1
                filter_row["mass2"] = sngl.mass2
                filter_row["spin1x"] = sngl.spin1x
                filter_row["spin1y"] = sngl.spin1y
                filter_row["spin1z"] = sngl.spin1z
                filter_row["spin2x"] = sngl.spin2x
                filter_row["spin2y"] = sngl.spin2y
                # FIXME should we keep this as seconds?
                # convert to nanoseconds to be consistent with all times in the database
                filter_row["template_duration"] = int(
                    sngl.template_duration * 1_000_000_000
                )
                self.filters["%04d" % int(bankid) + "_0"].append(filter_row)
        else:
            for i, subbank in enumerate(self.template_sngls):
                for template_id, sngl in subbank.items():
                    bankid = self.subbankids[i].split("_")[0]
                    filter_row = init_config_row(self.config["filter"])
                    filter_row["_filter_id"] = template_id
                    filter_row["bank_id"] = int(bankid)
                    filter_row["subbank_id"] = int(self.subbankids[i].split("_")[1])
                    filter_row["end_time_delta"] = (
                        sngl.end_time * 1_000_000_000 + sngl.end_time_ns
                    )
                    filter_row["mass1"] = sngl.mass1
                    filter_row["mass2"] = sngl.mass2
                    filter_row["spin1x"] = sngl.spin1x
                    filter_row["spin1y"] = sngl.spin1y
                    filter_row["spin1z"] = sngl.spin1z
                    filter_row["spin2x"] = sngl.spin2x
                    filter_row["spin2y"] = sngl.spin2y
                    filter_row["spin2z"] = sngl.spin2z
                    filter_row["template_duration"] = int(
                        sngl.template_duration * 1_000_000_000
                    )
                    self.filters["%04d" % int(bankid)].append(filter_row)

        #
        # Simulation
        #
        if self.injection_list:
            self.sims = []
            for inj in self.injection_list:
                sim_row = init_config_row(self.config["simulation"])
                sim_row["_simulation_id"] = inj.simulation_id
                sim_row["coa_phase"] = inj.coa_phase
                sim_row["distance"] = inj.distance
                sim_row["f_final"] = inj.f_final
                sim_row["f_lower"] = inj.f_lower
                sim_row["geocent_end_time"] = (
                    inj.geocent_end_time * 1_000_000_000 + inj.geocent_end_time_ns
                )
                sim_row["inclination"] = inj.inclination
                sim_row["polarization"] = inj.polarization
                sim_row["mass1"] = inj.mass1
                sim_row["mass2"] = inj.mass2
                sim_row["snr_H1"] = inj.alpha4
                sim_row["snr_L1"] = inj.alpha5
                sim_row["snr_V1"] = inj.alpha6
                sim_row["spin1x"] = inj.spin1x
                sim_row["spin1y"] = inj.spin1y
                sim_row["spin1z"] = inj.spin1z
                sim_row["spin2x"] = inj.spin2x
                sim_row["spin2y"] = inj.spin2y
                sim_row["spin2z"] = inj.spin2z
                sim_row["waveform"] = inj.waveform
                self.sims.append(sim_row)
        else:
            self.sims = None

        # Derived parameters needed by worker_process
        self.bankids = list(self.bankids)
        self.process = self.process_row
        self.process_params = self.params
        self.sim = self.sims
        self.config_segment = self.config["segment"]

        ParallelizeSinkElement.__post_init__(self)

        # FIXME Parallelize.enabled is only enabled once the
        # pipeline starts, so delay initializing the dbs until later
        self.init = False

    def get_username(self):
        try:
            return os.environ["LOGNAME"]
        except KeyError:
            pass
        try:
            return os.environ["USERNAME"]
        except KeyError:
            pass
        try:
            import pwd

            return pwd.getpwuid(os.getuid())[0]
        except (ImportError, KeyError):
            raise KeyError

    def process_outqueue(self, data):
        new_lr = data["new_lr"]
        frankenstein = data["frankenstein"]
        likelihood_ratio_upload = data["likelihood_ratio_upload"]
        bankid = data["bankid"]
        print(f"{asctime()} StillSuit Update dynamic", bankid, flush=True)
        self.strike_object.update_dynamic(
            bankid, frankenstein, likelihood_ratio_upload, new_lr
        )

    def get_state_from_queue(self):
        try:
            data = self.out_queue.get_nowait()
            self.process_outqueue(data)
        except Empty:
            return

    def pull(self, pad, frame):
        if frame.EOS:
            self.mark_eos(pad)

        pad_name = self.rsnks[pad]
        if pad_name == self.itacacac_pad_name:
            all_events = []
            all_triggers = []
            for buf_event in frame.events:
                if buf_event["trigger"] is None:
                    continue
                all_events.extend(buf_event["event"])
                all_triggers.extend(buf_event["trigger"])
            if all_events:
                data = {
                    "event_dict": {
                        "trigger": all_triggers,
                        "event": all_events,
                    }
                }
                self.in_queue.put(data)
        else:
            for buf in frame:
                if buf.data is not None:
                    data = {
                        "segment": {
                            self.segments_pad_map[pad_name]: (
                                frame.offset,
                                frame.end_offset,
                            )
                        }
                    }
                    self.in_queue.put(data)

    def internal(self):
        super().internal()
        if self.at_eos:
            for bankid in self.bankids:
                if self.is_online:
                    if self.injections:
                        fn_bankid = bankid + "_inj"
                    else:
                        fn_bankid = bankid + "_noninj"
                    desc = "%s_SGNL_TRIGGERS" % fn_bankid
                    fn = self.snapshot_filenames(desc)
                else:
                    # offline uses predefined output filenames
                    fn = self.trigger_output[bankid]
                sdict = {
                    "snapshot": {
                        "fn": fn,
                        "bankid": bankid,
                    }
                }
                self.in_queue.put(sdict)
            if self.terminated.is_set():
                print("At EOS and subprocess is terminated")
            else:
                drained_outq = self.sub_process_shutdown(600)
                print("after shutdown", len(drained_outq))
        else:
            if self.is_online:
                self.get_state_from_queue()
                for i, bankid in enumerate(self.bankids):
                    # FIXME: consider not looping over banks, since
                    # we have the subprocess
                    if self.injections:
                        fn_bankid = bankid + "_inj"
                    else:
                        fn_bankid = bankid + "_noninj"
                    desc = "%s_SGNL_TRIGGERS" % fn_bankid
                    if self.snapshot_ready(desc):
                        fn = self.snapshot_filenames(desc)
                        sdict = {
                            "snapshot": {
                                "fn": fn,
                                "bankid": bankid,
                            }
                        }
                        if self.injections:
                            if i == 0:
                                self.strike_object.load_rank_stat_pdf()
                            in_lr_file = self.strike_object.input_likelihood_file[
                                bankid
                            ]
                            sdict["snapshot"]["in_lr_file"] = in_lr_file
                        self.in_queue.put(sdict)

    def worker_process(
        self,
        context: WorkerContext,
        ifos: list,
        config_name: str,
        bankids: list,
        process: dict,
        process_params: list,
        sim: list,
        config_segment: dict,
        filters: dict,
    ):
        # Initialize worker state if needed
        if not context.state.get("subproc_init", False):
            context.state["dbs"], context.state["temp_segments"] = init_dbs(
                ifos,
                config_name,
                bankids,
                process,
                process_params,
                sim,
                filters,
            )
            context.state["subproc_init"] = True

        try:
            data = context.input_queue.get(timeout=1)
            if "segment" in data:
                append_segment(data, context.state["temp_segments"])
            elif "event_dict" in data:
                insert_event(data, context.state["dbs"])
            elif "snapshot" in data:
                lr_dict = on_snapshot(
                    data,
                    context.state["temp_segments"],
                    context.state["dbs"],
                    ifos,
                    config_name,
                    config_segment,
                    process,
                    process_params,
                    filters,
                    sim,
                    context.should_shutdown(),
                )
                if lr_dict is not None:
                    context.output_queue.put(lr_dict)
            else:
                raise ValueError("Unknown data")
        except Empty:
            return