madman.utilities

  1r"""Utilities."""
  2
  3import os
  4import sys
  5from collections.abc import Sequence
  6from contextlib import contextmanager
  7from math import ceil
  8from numbers import Real
  9from typing import Optional
 10
 11import numpy as np
 12from ase.parallel import paropen
 13
 14
 15def process_path(path: Optional[str]) -> Optional[str]:
 16    r"""Process path, creating tree if needed.
 17
 18    Args:
 19        path: Path to process.
 20
 21    Returns:
 22        Processed path.
 23    """
 24    if path is None:
 25        return None
 26    path = str(path)
 27    pathd = os.path.dirname(path)
 28    if pathd != "":
 29        os.makedirs(pathd, exist_ok=True)
 30    return path
 31
 32
 33@contextmanager
 34def redirect_to(prefix: Optional[str], *, mode: str = "w") -> None:
 35    r"""Redirect standard output/error to file.
 36
 37    Args:
 38        prefix: Prefix for output and error files.
 39        mode: Writing mode.
 40    """
 41    stdout_old = sys.stdout
 42    stderr_old = sys.stderr
 43
 44    prefix = process_path(prefix)
 45    if prefix is not None:
 46        sys.stdout = paropen(f"{prefix}.out", mode)
 47        sys.stderr = paropen(f"{prefix}.err", mode)
 48
 49    try:
 50        yield
 51
 52    finally:
 53        sys.stdout = stdout_old
 54        sys.stderr = stderr_old
 55
 56
 57def gen_regular_grid(xa: Real, xb: Real, dx: Real) -> np.ndarray[float]:
 58    """Generate regular grid.
 59
 60    Args:
 61        xa: Minimum value.
 62        xb: Maximum value.
 63        dx: Target increment.
 64
 65    Returns:
 66        Regular grid.
 67    """
 68    ndx = (xb - xa) / dx
 69    nx = ceil(ndx) + (1 if ndx % 1 == 0.0 else 0)
 70    return np.linspace(xa, xb, nx)
 71
 72
 73def correlate1d(
 74    x: Sequence[Real, ...],
 75    y1: Sequence[Real, ...],
 76    y2: Sequence[Real, ...],
 77) -> np.ndarray[float]:
 78    r"""Correlate two 1d signals.
 79
 80    Args:
 81        x: Independent variable values.
 82        y1: First signal dependent variable values.
 83        y2: Second signal dependent variable values.
 84
 85    Returns:
 86        Correlation independent and dependent variable values.
 87
 88    Raises:
 89        ValueError: If independent variable is not equally spaced.
 90        ValueError: If independent variable has duplicate values.
 91    """
 92    x = np.asarray(x, dtype=float).flatten()
 93    y1 = np.asarray(y1, dtype=float).flatten()
 94    y2 = np.asarray(y2, dtype=float).flatten()
 95
 96    i = np.argsort(x)
 97    x = x[i]
 98    y1 = y1[i]
 99    y2 = y2[i]
100
101    if not np.array_equal(x, np.unique(x)):
102        raise ValueError("Duplicate independent variable values!")
103
104    dx = np.diff(x)
105    xn = len(dx)
106    dxm = np.mean(dx)
107    if not np.all(np.isclose(dx, dxm)):
108        raise ValueError("Independent variable not equally spaced!")
109
110    x = dxm * np.arange(-xn, xn + 1)
111    z = dxm * np.correlate(y2, y1, mode="full")
112    return x, z