pytesprocess 0.1.1__py3-none-any.whl

pytesprocess/salting/sampling.py

@@ -0,0 +1,84 @@
+"""Sampling helpers for salting metadata generation."""
+
+from __future__ import annotations
+
+import cloudpickle
+import numpy as np
+
+
+def sample_pdf(function, xrange, nsamples=1000, npoints=10000,
+               normalize_cdf=True, rng=None):
+    """Sample a one-dimensional PDF with inverse-transform sampling."""
+    if not callable(function):
+        raise TypeError("Input PDF must be callable.")
+    if len(xrange) != 2 or float(xrange[1]) <= float(xrange[0]):
+        raise ValueError("xrange must contain [xmin, xmax] with xmax > xmin.")
+    if int(nsamples) < 0:
+        raise ValueError("nsamples must be non-negative.")
+    if int(npoints) < 2:
+        raise ValueError("npoints must be at least 2.")
+
+    x = np.linspace(float(xrange[0]), float(xrange[1]), num=int(npoints))
+    pdf = np.asarray(function(x), dtype=float)
+    if pdf.shape != x.shape:
+        raise ValueError("PDF function must return one value per input point.")
+    if np.any(~np.isfinite(pdf)) or np.any(pdf < 0):
+        raise ValueError("PDF values must be finite and non-negative.")
+
+    # Cumulative trapezoidal integral without requiring scipy.  This works on
+    # NumPy 1.26 and current NumPy releases.
+    dx = np.diff(x)
+    areas = 0.5 * (pdf[:-1] + pdf[1:]) * dx
+    cdf = np.concatenate(([0.0], np.cumsum(areas)))
+    total = float(cdf[-1])
+    if total <= 0:
+        raise ValueError("PDF has zero integral over the requested range.")
+    if normalize_cdf:
+        cdf = cdf / total
+
+    # np.interp requires monotonically increasing sample points. Flat regions
+    # of a PDF create duplicate CDF values, so retain the first occurrence.
+    cdf_unique, unique_inds = np.unique(cdf, return_index=True)
+    x_unique = x[unique_inds]
+    if len(cdf_unique) < 2:
+        raise ValueError("Unable to construct an invertible CDF from the PDF.")
+
+    if rng is None:
+        rng = np.random.default_rng()
+    high = 1.0 if normalize_cdf else total
+    samples = rng.random(int(nsamples)) * high
+    return np.interp(samples, cdf_unique, x_unique)
+
+
+def sample_dm_distributions(pdf_file, nsamples_per_mass, *,
+                            energy_unit_to_ev=1.0e3,
+                            energy_range=(1.0e-5, 1.0),
+                            random_seed=None):
+    """Sample the historical cloudpickle DM distribution file.
+
+    Returns
+    -------
+    energies_ev : ndarray
+        Recoil energies in eV.
+    masses_mev : ndarray
+        DM mass corresponding to each sampled recoil energy.
+    """
+    rng = np.random.default_rng(random_seed)
+    with open(pdf_file, "rb") as handle:
+        distributions = cloudpickle.load(handle)
+
+    energies = []
+    masses = []
+    for mass, data in distributions.items():
+        if not isinstance(data, dict) or "dmrate" not in data:
+            raise ValueError(
+                f'DM distribution entry for mass {mass!r} is missing "dmrate".'
+            )
+        sampled = sample_pdf(
+            data["dmrate"], energy_range,
+            nsamples=int(nsamples_per_mass), rng=rng,
+        )
+        energies.extend(np.asarray(sampled) * float(energy_unit_to_ev))
+        masses.extend([mass] * len(sampled))
+
+    return np.asarray(energies, dtype=float), np.asarray(masses)

pytesprocess/utils/__init__.py

@@ -0,0 +1,5 @@
+from .utils import *
+from .arg_utils import *
+from .filter_hdf5 import *
+
+from .dataframe_output import *

pytesprocess/utils/arg_utils.py

@@ -0,0 +1,122 @@
+import re
+
+import numpy as np
+
+
+def build_range_str(data_list):
+    """Build a compact range string such as ``1-3_7_9-11``."""
+    data_list.sort()
+    units = []
+    prev_val = data_list[0]
+
+    for val in data_list:
+        if val == prev_val + 1:
+            units[-1].append(val)
+        else:
+            units.append([val])
+        prev_val = val
+
+    return "_".join(
+        f"{unit[0]}-{unit[-1]}" if len(unit) > 1 else str(unit[0])
+        for unit in units
+    )
+
+
+def hyphen_range(data_str):
+    """Expand strings such as ``1-3,7,9-11`` into a sorted integer list."""
+    data_str = "".join(data_str.split())
+    data_str = data_str.replace(",", "_")
+    values = set()
+
+    for item in data_str.split("_"):
+        bounds = item.split("-")
+        if len(bounds) not in (1, 2):
+            raise SyntaxError(
+                "hyphen_range received an incorrectly formatted argument: "
+                + data_str.replace("_", ",")
+            )
+        if len(bounds) == 1:
+            values.add(int(bounds[0]))
+        else:
+            values.update(range(int(bounds[0]), int(bounds[1]) + 1))
+
+    return sorted(values)
+
+
+def extract_list(arg_list):
+    """Extract a list separated by whitespace, comma, or semicolon."""
+    if not isinstance(arg_list, list):
+        arg_list = [arg_list]
+
+    output_list = []
+    for item in arg_list:
+        output_list.extend(
+            token for token in re.split(r"[\s,;]+", item.strip()) if token
+        )
+    return output_list
+
+
+def convert_to_seconds(par):
+    """Convert a duration string to seconds."""
+    par = str(par)
+    if par.isdigit():
+        return float(par)
+
+    seconds_per_unit = {
+        "us": 1e-6,
+        "ms": 1e-3,
+        "s": 1,
+        "m": 60,
+        "h": 3600,
+        "d": 86400,
+        "w": 604800,
+    }
+
+    suffix_length = 2 if "us" in par or "ms" in par else 1
+    unit = par[-suffix_length:]
+    if unit not in seconds_per_unit:
+        raise ValueError(f'ERROR: time unit "{unit}" unknown!')
+
+    return float(par[:-suffix_length]) * seconds_per_unit[unit]
+
+
+def extract_stream_id(stream):
+    if isinstance(stream, str):
+        if stream.startswith("I"):
+            return stream
+        if not stream.isdigit():
+            stream = extract_stream_num(stream)
+
+    stream_str = str(stream)
+    return (
+        "I"
+        + stream_str[:-14]
+        + "_D"
+        + stream_str[-14:-6]
+        + "_T"
+        + stream_str[-6:]
+    )
+
+
+def extract_stream_num(stream):
+    if isinstance(stream, int):
+        return stream
+    if not isinstance(stream, str):
+        raise ValueError("ERROR in extract_stream_num: stream should be a string")
+
+    stream_split = stream.split("_")
+    if (
+        len(stream_split) < 3
+        or not stream_split[-1].startswith("T")
+        or not stream_split[-2].startswith("D")
+        or not stream_split[-3].startswith("I")
+    ):
+        raise ValueError("ERROR in extract_stream_num: unknown stream name format")
+
+    return np.uint64(
+        float(
+            stream_split[-3][1:]
+            + stream_split[-2][1:]
+            + stream_split[-1][1:]
+        )
+    )

pytesprocess/utils/dataframe_output.py

@@ -0,0 +1,120 @@
+"""Common naming and identity helpers for processed dataframe products."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+import os
+import stat
+
+from .arg_utils import extract_stream_id, extract_stream_num
+
+
+@dataclass(frozen=True)
+class DataframeGroup:
+    """Identity and filesystem location of one processed dataframe product."""
+
+    dataframe_type: str
+    name: str
+    group_id: str
+    group_number: int
+    path: str
+    processing_label: str | None = None
+
+    def file_path(self, file_index: int) -> str:
+        index = int(file_index)
+        if index < 1:
+            raise ValueError("dataframe_file_index must be >= 1")
+        return str(Path(self.path) / f"{self.name}_F{index:04d}.hdf5")
+
+    def row_metadata(self, file_index: int) -> dict:
+        return {
+            "dataframe_type": self.dataframe_type,
+            "dataframe_group_name": self.name,
+            "dataframe_group_id": self.group_id,
+            "dataframe_group_number": self.group_number,
+            "dataframe_file_index": int(file_index),
+            "processing_label": self.processing_label,
+        }
+
+
+def create_dataframe_group(
+    base_path,
+    *,
+    dataframe_type,
+    facility,
+    processing_label=None,
+    restricted=False,
+    data_type="background",
+    output_group_name=None,
+    name_prefix=None,
+    now=None,
+):
+    """Create one processed-dataframe group with one canonical ``I#_D#_T#`` ID.
+
+    ``F####`` is intentionally not part of the group identity.  It is a shard
+    index assigned to an independent processing task/worker.
+    """
+
+    dataframe_type = str(dataframe_type).strip().lower()
+    if not dataframe_type:
+        raise ValueError("dataframe_type is required")
+
+    base_path = Path(base_path)
+    now = datetime.now() if now is None else now
+
+    if output_group_name is None:
+        group_id = f"I{int(facility)}_D{now:%Y%m%d}_T{now:%H%M%S}"
+        prefix = str(name_prefix).strip() if name_prefix else dataframe_type
+        if processing_label:
+            prefix = f"{processing_label}_{prefix}"
+        if restricted:
+            prefix += "_restricted"
+        elif data_type == "calibration":
+            prefix += "_calibration"
+        group_name = f"{prefix}_{group_id}"
+        output_dir = base_path / group_name
+    else:
+        group_name = str(output_group_name)
+        try:
+            group_id = extract_stream_id(group_name)
+        except Exception as exc:
+            raise ValueError(
+                "output_group_name must contain a canonical I#_Dyyyymmdd_Thhmmss ID"
+            ) from exc
+        output_dir = base_path if base_path.name == group_name else base_path / group_name
+
+    group_number = int(extract_stream_num(group_id))
+    output_dir.mkdir(parents=True, exist_ok=True)
+    try:
+        os.chmod(
+            output_dir,
+            stat.S_IRWXG | stat.S_IRWXU | stat.S_IROTH | stat.S_IXOTH,
+        )
+    except OSError:
+        # Directory creation is the important operation; chmod can legitimately
+        # fail on some shared/object-backed filesystems.
+        pass
+
+    return DataframeGroup(
+        dataframe_type=dataframe_type,
+        name=group_name,
+        group_id=group_id,
+        group_number=group_number,
+        path=str(output_dir),
+        processing_label=processing_label,
+    )
+
+
+def add_dataframe_group_columns(dataframe, group: DataframeGroup, file_index: int):
+    """Attach canonical dataframe-group identity columns to a Vaex dataframe."""
+    import numpy as np
+
+    count = len(dataframe)
+    for key, value in group.row_metadata(file_index).items():
+        if value is None:
+            dataframe[key] = np.asarray([None] * count, dtype=object)
+        else:
+            dataframe[key] = np.asarray([value] * count)
+    return dataframe