sdatip 1.0.0__py3-none-any.whl

sdatip/__init__.py

@@ -0,0 +1,46 @@
+"""SDATIP: Fast Stochastic Determination of Arrival Time and Initial Polarity.
+
+A high-performance Python package for seismic waveform analysis that determines
+arrival time and initial polarity using Markov chain-based stochastic methods.
+
+Example usage:
+    >>> import sdatip
+    >>> 
+    >>> # Quick processing
+    >>> result = sdatip.process_waveform(
+    ...     name="station_A",
+    ...     data=waveform_array,
+    ...     output_dir="./output/"
+    ... )
+    >>> print(f"Arrival time: {result['arrival_time']:.3f}s")
+    >>> print(f"Polarity: {result['polarity_up']:.3f}")
+    >>> 
+    >>> # Step-by-step processing
+    >>> wf = sdatip.Waveform("station_A")
+    >>> wf.importdata(data, delta=0.01)
+    >>> wf.analyzedata()
+    >>> # ... continue with other steps
+"""
+
+from sdatip.waveform import Waveform, findnoise
+from sdatip.state import State
+from sdatip.pmi import entropy, pmi, maxpmi, calculate_general
+from sdatip.processor import process_waveform, process_batch
+
+__version__ = "1.0.0"
+__author__ = "Chuan1937"
+
+__all__ = [
+    # Main classes
+    "Waveform",
+    "State",
+    # High-level API
+    "process_waveform",
+    "process_batch",
+    # Utility functions
+    "findnoise",
+    "entropy",
+    "pmi",
+    "maxpmi",
+    "calculate_general",
+]

sdatip/plotting.py

@@ -0,0 +1,404 @@
+"""Plotting utilities for seismic waveform analysis results."""
+
+import numpy as np
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+
+def plot_result(wf, state, qualifiedid, name, outputdir):
+    """Generate detailed probability plot for analysis results.
+
+    Args:
+        wf: Waveform object with processed data.
+        state: State object with probability estimates.
+        qualifiedid: Index of the solution to plot.
+        name: Station name for filename.
+        outputdir: Directory to save the plot.
+    """
+    plt.rcParams.update({
+        "font.weight": "bold",
+        "axes.labelweight": "bold",
+        "font.family": "Times New Roman",
+        "font.size": 45,
+    })
+
+    fig = plt.figure(figsize=(25, 9))
+    timeprob = state.timeprob[qualifiedid]
+
+    b_upthreshold = np.array(state.upthreshold)
+    b_downthreshold = np.array(state.downthreshold)
+    b_Apeak = np.array([item[0] for item in state.Apeak])
+
+    a_cut = np.array(wf.cut).flatten().astype(int)
+
+    b_upthreshold[-1] = b_upthreshold[-2] * 1.5
+
+    prob1 = timeprob / (b_upthreshold - b_downthreshold)
+    alphacoefficient = (0.75 - 0.03) / np.max(prob1)
+    alphas = 0.03 + prob1 * alphacoefficient
+
+    colori = np.zeros((state.num, 3))
+    colori[b_Apeak > 0] = [1, 0, 0]
+    colori[b_Apeak < 0] = [0, 1, 0]
+
+    mask_nonzero = np.abs(b_Apeak) > 0
+
+    downt = np.zeros(state.num)
+    upt = np.zeros(state.num)
+
+    cut_indices = a_cut[mask_nonzero]
+    tchange = wf.longtimestamp[cut_indices + 1] - wf.longtimestamp[cut_indices]
+    achange = wf.denselongdata[cut_indices + 1] - wf.denselongdata[cut_indices]
+    achange[achange == 0] = 1e-9
+
+    downt[mask_nonzero] = (
+        tchange / achange * (b_downthreshold[mask_nonzero] - np.abs(wf.denselongdata[cut_indices]))
+        + wf.longtimestamp[cut_indices]
+    )
+    upt[mask_nonzero] = (
+        tchange / achange * (b_upthreshold[mask_nonzero] - np.abs(wf.denselongdata[cut_indices]))
+        + wf.longtimestamp[cut_indices]
+    )
+
+    downt[~mask_nonzero] = wf.longtimestamp[-1]
+    upt[~mask_nonzero] = wf.longtimestamp[-1] + 0.1 * wf.timestamp[-1]
+
+    ax1 = fig.add_axes([0.4, 0.33, 0.55, 0.6])
+    ax1.plot(wf.longtimestamp, wf.denselongdata, linewidth=2.5, color="k", linestyle="-")
+    ax1.plot(wf.longtimestamp, abs(wf.denselongdata), linewidth=2.5, color="k", linestyle=":", alpha=0.9)
+
+    ylim_min = -1 * b_upthreshold[-1]
+    ax1.vlines(
+        downt[mask_nonzero], ylim_min, b_downthreshold[mask_nonzero],
+        color="k", linestyle="--", alpha=0.3, linewidth=0.5
+    )
+    ax1.vlines(
+        upt[mask_nonzero], ylim_min, b_upthreshold[mask_nonzero],
+        color="k", linestyle="--", alpha=0.3, linewidth=0.5
+    )
+    ax1.hlines(
+        b_downthreshold[mask_nonzero], 0, downt[mask_nonzero],
+        color="k", linestyle="--", alpha=0.3, linewidth=0.5
+    )
+    ax1.hlines(
+        b_upthreshold[mask_nonzero], 0, upt[mask_nonzero],
+        color="k", linestyle="--", alpha=0.3, linewidth=0.5
+    )
+
+    for i in range(state.num):
+        ax1.fill_between(
+            [0, downt[i], upt[i]],
+            [b_downthreshold[i], b_downthreshold[i], b_upthreshold[i]],
+            [b_upthreshold[i], b_upthreshold[i], b_upthreshold[i]],
+            color=colori[i], alpha=alphas[i]
+        )
+        ax1.fill_betweenx(
+            [-1 * b_upthreshold[-2], b_downthreshold[i], b_upthreshold[i]],
+            [downt[i], downt[i], upt[i]],
+            [upt[i], upt[i], upt[i]],
+            color=colori[i], alpha=alphas[i]
+        )
+
+    ax1.set(
+        xlim=(0, np.max(wf.densetimestamp) + 0.1 * wf.timestamp[-1]),
+        ylim=(-1 * b_upthreshold[-2], b_upthreshold[-1])
+    )
+    ax1.tick_params(direction="out", size=20)
+    ax1.set_yticks(np.array([-1 * b_upthreshold[-2], 0, b_upthreshold[-1]]))
+    ax1.set_yticklabels(
+        ["%.2f" % (-1 * b_upthreshold[-2]), "%.2f" % 0, "%.2f" % b_upthreshold[-1]],
+        fontweight="bold"
+    )
+    ax1.set_xticks(np.linspace(0, wf.timestamp[-1], 5))
+    ax1.set_xticklabels(
+        ["%.2f" % (0), "%.2f" % (wf.timestamp[-1] / 4), "%.2f" % (wf.timestamp[-1] / 4 * 2),
+         "%.2f" % (wf.timestamp[-1] / 4 * 3), "%.2f" % (wf.timestamp[-1])],
+        fontweight="bold"
+    )
+
+    for spine in ax1.spines.values():
+        spine.set_linewidth(2)
+
+    ax2 = fig.add_axes([
+        0.1, 0.93 - 0.6 / (b_upthreshold[-1] + b_upthreshold[-2]) * b_upthreshold[-1],
+        0.25, 0.6 / (b_upthreshold[-1] + b_upthreshold[-2]) * b_upthreshold[-1]
+    ])
+    ax2.invert_xaxis()
+
+    ax2.barh(
+        y=b_downthreshold, width=prob1, height=b_upthreshold - b_downthreshold,
+        left=0, align="edge", color=colori, alpha=1
+    )
+
+    ax2.set(ylim=(0, b_upthreshold[-1]))
+    max_prob1 = np.max(prob1)
+    ax2.set_xticks(np.linspace(0, max_prob1, 5))
+    ax2.set_xticklabels(
+        ["%.2f" % (0), "%.2f" % (max_prob1 / 4), "%.2f" % (max_prob1 / 2),
+         "%.2f" % (max_prob1 * 3 / 4), "%.2f" % (max_prob1)],
+        fontweight="bold"
+    )
+    ax2.set_yticks(np.linspace(0, b_upthreshold[-1], 5))
+    ax2.set_yticklabels(
+        ["%.2f" % (0), "%.2f" % (b_upthreshold[-1] / 4), "%.2f" % (b_upthreshold[-1] / 2),
+         "%.2f" % (b_upthreshold[-1] * 3 / 4), "%.2f" % (b_upthreshold[-1])],
+        fontweight="bold"
+    )
+    ax2.tick_params(direction="out", size=20)
+    ax2.set_ylabel(r"$\mathbf{\varepsilon_{threshold}}$", weight="bold")
+    ax2.set_xlabel(r"PDF of $\mathbf{\varepsilon_{threshold}}$", weight="bold")
+    for spine in ax2.spines.values():
+        spine.set_linewidth(2)
+
+    ax3 = fig.add_axes([0.4, 0.1, 0.55, 0.15], sharex=ax1)
+
+    max_prob1_1 = np.max(prob1) * 1.1
+    ax3.vlines(downt, 0, max_prob1_1, color="k", linestyle="-", alpha=0.3, linewidth=0.5)
+    ax3.vlines(upt, 0, max_prob1_1, color="k", linestyle="-", alpha=0.3, linewidth=0.5)
+    ax3.hlines(prob1, downt, upt, color="k", linestyle="-", alpha=0.3, linewidth=0.5)
+    ax3.bar(x=downt, height=prob1, width=upt - downt, align="edge", color=colori, alpha=1)
+    ax3.plot(
+        [state.arrivalestimate, state.arrivalestimate], [0, max_prob1_1],
+        linewidth=2.3, color="k", linestyle=":", alpha=0.9
+    )
+
+    ax3.set(ylim=(0, max_prob1_1))
+    ax3.set_yticks(np.linspace(0, np.max(prob1), 3))
+    ax3.set_yticklabels(
+        ["%.2f" % (0), "%.2f" % (np.max(prob1) / 2), "%.2f" % (np.max(prob1))],
+        fontweight="bold"
+    )
+    ax3.set_xticks(np.linspace(0, wf.timestamp[-1], 5))
+    ax3.set_xticklabels(
+        ["%.2f" % (0), "%.2f" % (wf.timestamp[-1] / 4), "%.2f" % (wf.timestamp[-1] / 2),
+         "%.2f" % (wf.timestamp[-1] * 3 / 4), "%.2f" % (wf.timestamp[-1])],
+        fontweight="bold"
+    )
+    ax3.tick_params(direction="out", size=20)
+    ax3.set_ylabel("PDF of Time", weight="bold")
+    ax3.set_xlabel("Time/s", weight="bold")
+    for spine in ax3.spines.values():
+        spine.set_linewidth(2)
+
+    ax4 = fig.add_axes([0.1, 0.13, 0.23, 0.05])
+
+    width = [float(state.polarityup), float(state.polarityunknown), float(state.polaritydown)]
+    left = [0, width[0], width[0] + width[1]]
+    colors = [[1, 0, 0], [0.7, 0.7, 0.7], [0, 1, 0]]
+    labels = ["Up", "Unknown", "Down"]
+    ax4.barh(y=[1, 1, 1], width=width, height=1, left=left, color=colors)
+
+    ax4.set_xticks([0.5])
+    ax4.set_xticklabels(["0.5"], fontweight="bold")
+    ax4.text(0, 2, "Up:%.1f%%" % (width[0] * 100))
+    ax4.text(0.85, 2, "Down:%.1f%%" % (width[2] * 100))
+    ax4.set_yticks([])
+    ax4.set(xlim=(0, 1), ylim=(0.5, 1.5))
+    ax4.plot([0.5, 0.5], [0.5, 1.5], linewidth=2.5, color="k", linestyle=":", alpha=1)
+    for spine in ax4.spines.values():
+        spine.set_linewidth(2)
+    ax4.legend(
+        handles=[plt.Rectangle((0, 0), 1, 1, color=c) for c in colors],
+        labels=labels, ncol=3, loc="lower center", bbox_to_anchor=(0.5, 1.5)
+    )
+
+    fig.text(0.21, 0.38, "%s" % (name), {"fontweight": "bold", "fontsize": 25}, horizontalalignment="center")
+    fig.text(0.24, 0.33, r"$\mathbf{A_{peak}}$" + ": %.3f" % (state.Apeakestimate), {"fontweight": "bold", "fontsize": 15})
+    fig.text(0.24, 0.28, r"$\mathbf{\sigma}$" + ": %.3f" % (state.sigmaestimate), {"fontweight": "bold", "fontsize": 15})
+    fig.text(0.1, 0.33, "Arrivaltime" + ": %.3f" % (state.arrivalestimate), {"fontweight": "bold", "fontsize": 15})
+    fig.text(0.1, 0.28, "Polarity Up" + ": %.3f" % (state.polarityestimation), {"fontweight": "bold", "fontsize": 15})
+    fig.text(0.1, 0.23, "Eig value:" + " %s" % (state.bigeig))
+
+    fig.savefig("%s" % (outputdir) + "%s_%d.pdf" % (name, qualifiedid))
+    plt.close(fig)
+
+
+def plot_result_graduate(wf, state, qualifiedid, name, outputdir):
+    """Generate graduate-style plot for publication.
+
+    Args:
+        wf: Waveform object with processed data.
+        state: State object with probability estimates.
+        qualifiedid: Index of the solution to plot.
+        name: Station name for filename.
+        outputdir: Directory to save the plot.
+    """
+    plt.rcParams.update({
+        "font.weight": "normal",
+        "axes.labelweight": "normal",
+        "font.family": "Times New Roman",
+        "font.size": 45,
+    })
+
+    fig = plt.figure(figsize=(20, 12))
+    timeprob = state.timeprob[qualifiedid]
+
+    b_upthreshold = np.array(state.upthreshold)
+    b_downthreshold = np.array(state.downthreshold)
+    b_Apeak = np.array([item[0] for item in state.Apeak])
+
+    a_cut = np.array(wf.cut).flatten().astype(int)
+
+    b_upthreshold[-1] = b_upthreshold[-2] * 1.5
+
+    prob1 = timeprob / (b_upthreshold - b_downthreshold)
+    alphacoefficient = 0.75 / np.max(prob1)
+    alphas = timeprob / (b_upthreshold - b_downthreshold) * alphacoefficient
+    alphas_fill = 0.03 + alphas
+
+    colori = np.zeros((state.num, 3))
+    colori[b_Apeak > 0] = [1, 0, 0]
+    colori[b_Apeak < 0] = [0, 0, 1]
+
+    oritprob = np.zeros(wf.length)
+
+    tprobid = np.floor(wf.longtimestamp[a_cut] / wf.delta).astype(int)
+    np.add.at(oritprob, tprobid, timeprob)
+
+    mask_nonzero = np.abs(b_Apeak) > 0
+
+    downt1 = np.zeros(state.num)
+    upt1 = np.zeros(state.num)
+
+    downt1[mask_nonzero] = wf.timestamp[tprobid[mask_nonzero]]
+    upt1[mask_nonzero] = wf.timestamp[tprobid[mask_nonzero] + 1]
+    downt1[~mask_nonzero] = wf.timestamp[-1]
+    upt1[~mask_nonzero] = wf.timestamp[-1] + 0.1 * wf.timestamp[-1]
+
+    downt = np.zeros(state.num)
+    upt = np.zeros(state.num)
+
+    cut_indices = a_cut[mask_nonzero]
+    tchange = wf.longtimestamp[cut_indices + 1] - wf.longtimestamp[cut_indices]
+    achange = wf.denselongdata[cut_indices + 1] - wf.denselongdata[cut_indices]
+    achange[achange == 0] = 1e-9
+
+    downt[mask_nonzero] = (
+        tchange / achange * (b_downthreshold[mask_nonzero] - np.abs(wf.denselongdata[cut_indices]))
+        + wf.longtimestamp[cut_indices]
+    )
+    upt[mask_nonzero] = (
+        tchange / achange * (b_upthreshold[mask_nonzero] - np.abs(wf.denselongdata[cut_indices]))
+        + wf.longtimestamp[cut_indices]
+    )
+
+    downt[~mask_nonzero] = wf.longtimestamp[-1]
+    upt[~mask_nonzero] = wf.longtimestamp[-1] + 0.1 * wf.timestamp[-1]
+
+    ax1 = fig.add_axes([0.47, 0.33, 0.45, 0.6])
+    ax1.plot(wf.longtimestamp, wf.denselongdata, linewidth=5, color="k", linestyle="-")
+    ax1.plot(wf.longtimestamp, abs(wf.denselongdata), linewidth=5, color="k", linestyle=":", alpha=0.9)
+
+    for i in range(state.num):
+        ax1.fill_between(
+            [0, downt[i], upt[i]],
+            [b_downthreshold[i], b_downthreshold[i], b_upthreshold[i]],
+            [b_upthreshold[i], b_upthreshold[i], b_upthreshold[i]],
+            color=colori[i], alpha=alphas_fill[i]
+        )
+        ax1.fill_betweenx(
+            [-1 * b_upthreshold[-2], b_downthreshold[i], b_upthreshold[i]],
+            [downt[i], downt[i], upt[i]],
+            [upt[i], upt[i], upt[i]],
+            color=colori[i], alpha=alphas[i]
+        )
+
+    ax1.set(
+        xlim=(0, np.max(wf.densetimestamp) + 0.1 * wf.timestamp[-1]),
+        ylim=(-1 * b_upthreshold[-2], b_upthreshold[-1])
+    )
+    ax1.yaxis.tick_right()
+    ax1.tick_params(direction="out", size=20, length=5, width=2)
+    ax1.set_yticks(np.array([-1 * b_upthreshold[-2], 0, b_upthreshold[-1]]))
+    ax1.set_yticklabels([
+        "%.1f" % (-1 * b_upthreshold[-2] / 10000),
+        "%.1f" % (0),
+        "%.1f" % (b_upthreshold[-1] / 10000)
+    ])
+    ax1.set_xticks(np.linspace(0, wf.timestamp[-1], 5))
+    ax1.set_xticklabels(
+        ["%.2f" % (0), "%.2f" % (wf.timestamp[-1] / 4), "%.2f" % (wf.timestamp[-1] / 4 * 2),
+         "%.2f" % (wf.timestamp[-1] / 4 * 3), "%.2f" % (wf.timestamp[-1])],
+        fontsize=25
+    )
+    [t.set_color("white") for t in ax1.xaxis.get_ticklabels()]
+    for spine in ax1.spines.values():
+        spine.set_linewidth(5)
+
+    ax2 = fig.add_axes([
+        0.11, 0.93 - 0.6 / (b_upthreshold[-1] + b_upthreshold[-2]) * b_upthreshold[-1],
+        0.25, 0.6 / (b_upthreshold[-1] + b_upthreshold[-2]) * b_upthreshold[-1]
+    ])
+    ax2.invert_xaxis()
+    ax2.barh(
+        y=b_downthreshold, width=prob1, height=b_upthreshold - b_downthreshold,
+        left=0, align="edge", color=colori, alpha=1
+    )
+
+    ax2.set(ylim=(-4000, b_upthreshold[-1]))
+    max_prob1 = np.max(prob1)
+    ax2.set_xticks(np.linspace(0, max_prob1, 3))
+    ax2.set_xticklabels(["%d" % (0), "%.2f" % (max_prob1 / 2 * 70), "%.2f" % (max_prob1 * 70)])
+    ax2.set_yticks(np.linspace(0, b_upthreshold[-1], 3))
+    ax2.set_yticklabels([
+        "%.1f" % (0),
+        "%.1f" % (b_upthreshold[-1] / 2 / 10000),
+        "%.1f" % (b_upthreshold[-1] / 10000)
+    ])
+    ax2.tick_params(direction="out", size=20, length=5, width=2)
+    ax2.set_ylabel(r"$\epsilon$")
+    ax2.set_xlabel("PDF")
+    for spine in ax2.spines.values():
+        spine.set_linewidth(5)
+
+    ax3 = fig.add_axes([0.47, 0.13, 0.45, 0.12])
+    ax3.bar(x=downt1, height=oritprob[tprobid], width=upt1 - downt1, align="edge", color=colori, alpha=1)
+
+    max_oritprob = np.max(oritprob)
+    ax3.plot(
+        [state.arrivalestimate, state.arrivalestimate], [0, max_oritprob * 1.1],
+        linewidth=3, color="k", linestyle=":", alpha=0.9
+    )
+    ax3.yaxis.tick_right()
+    ax3.set(
+        xlim=(0, np.max(wf.densetimestamp) + 0.1 * wf.timestamp[-1]),
+        ylim=(0, max_oritprob * 1.1)
+    )
+    ax3.set_yticks(np.linspace(0, max_oritprob, 3))
+    ax3.set_yticklabels(["%d" % (0), "%d" % (100 * int(max_oritprob / 2)), "%d" % (100 * int(max_oritprob))])
+    ax3.set_xticks(np.linspace(0, wf.timestamp[-1], 5))
+    ax3.set_xticklabels(
+        ["%.2f" % (0), "%.2f" % (wf.timestamp[-1] / 4), "%.2f" % (wf.timestamp[-1] / 4 * 2),
+         "%.2f" % (wf.timestamp[-1] / 4 * 3), "%.2f" % (wf.timestamp[-1])],
+        fontsize=45
+    )
+    ax3.tick_params(direction="out", size=20, length=5, width=1)
+    ax3.set_ylabel("PDF")
+    ax3.set_xlabel("Time (s)")
+    for spine in ax3.spines.values():
+        spine.set_linewidth(5)
+
+    ax4 = fig.add_axes([0.125, 0.12, 0.225, 0.05])
+    width = [float(state.polarityup), float(state.polarityunknown), float(state.polaritydown)]
+    left = [0, width[0], width[0] + width[1]]
+    colors = [[1, 0, 0], [0.7, 0.7, 0.7], [0, 0, 1]]
+    ax4.barh(y=[1, 1, 1], width=width, height=1, left=left, color=colors)
+
+    ax4.set_xticks([0.5])
+    ax4.set_xticklabels(["0.5"])
+    ax4.text(-0.07, 2, "U: %.1f%%" % (abs(width[0]) * 100), fontsize=45)
+    ax4.text(0.7, 2, "D: %.1f%%" % (abs(width[2]) * 100), fontsize=45)
+    ax4.text(-0.37, 2, "Pol:", fontsize=45)
+    ax4.set_yticks([])
+    ax4.set(xlim=(0, 1), ylim=(0.5, 1.5))
+    ax4.plot([0.5, 0.5], [0.5, 1.5], linewidth=5, color="k", linestyle=":", alpha=1)
+    for spine in ax4.spines.values():
+        spine.set_linewidth(5)
+
+    fig.text(0.04, 0.305, f"E.V.: {state.bigeig[0]:.0f}, {state.bigeig[1]:.0f}, {state.bigeig[2]:.2f}, ...", fontsize=45)
+
+    fig.savefig("%s" % (outputdir) + "%s.eps" % (name))
+    fig.savefig("%s" % (outputdir) + "%s_%d.pdf" % (name, qualifiedid))
+    plt.close(fig)

sdatip/pmi.py

@@ -0,0 +1,155 @@
+"""Pointwise Mutual Information (PMI) calculations for seismic waveform analysis.
+
+This module provides JIT-compiled functions for computing PMI and entropy
+to find optimal cut points in binary sequences.
+"""
+
+import numba
+import numpy as np
+
+
+@numba.jit(nopython=True, cache=True)
+def entropy(ampen2):
+    """Calculate binary entropy of a 0/1 sequence.
+
+    Uses Shannon entropy formula: H = -p0*log2(p0) - p1*log2(p1)
+    """
+    if ampen2.size == 0:
+        return 0.0
+    counts = np.bincount(ampen2)
+    if len(counts) < 2:
+        return 0.0
+    zero_count = counts[0]
+    one_count = counts[1]
+    if zero_count == 0 or one_count == 0:
+        return 0.0
+    length = len(ampen2)
+    p_zero = zero_count / length
+    p_one = one_count / length
+    return -(np.log2(p_zero) * p_zero + np.log2(p_one) * p_one)
+
+
+@numba.jit(nopython=True, cache=True)
+def pmi(ampbi, t):
+    """Calculate Pointwise Mutual Information at split point t.
+
+    Computes PMI matrix between two parts of binary sequence split at index t.
+    Returns (mi_value, normalized_pmi_value).
+    """
+    length_all = len(ampbi)
+    part1 = ampbi[:t]
+    part2 = ampbi[t:]
+
+    zero_before_split = (part1 == 0).sum()
+    one_before_split = (part1 == 1).sum()
+    zero_after_split = (part2 == 0).sum()
+    one_after_split = (part2 == 1).sum()
+
+    pmi_matrix = np.zeros((2, 2))
+    npmi_matrix = np.zeros((2, 2))
+
+    total_zeros = zero_before_split + zero_after_split
+    total_ones = one_before_split + one_after_split
+    len_part1 = zero_before_split + one_before_split
+    len_part2 = zero_after_split + one_after_split
+
+    if zero_before_split > 0 and total_zeros > 0 and len_part1 > 0:
+        denominator = len_part1 * total_zeros
+        pmi_val = zero_before_split * length_all / denominator
+        pmi_matrix[0, 0] = (zero_before_split / length_all) * np.log2(pmi_val)
+        npmi_matrix[0, 0] = -np.log2(pmi_val) / np.log2(zero_before_split / length_all)
+    else:
+        npmi_matrix[0, 0] = -1.0
+
+    if one_before_split > 0 and total_ones > 0 and len_part1 > 0:
+        denominator = len_part1 * total_ones
+        pmi_val = one_before_split * length_all / denominator
+        pmi_matrix[0, 1] = (one_before_split / length_all) * np.log2(pmi_val)
+        npmi_matrix[0, 1] = -np.log2(pmi_val) / np.log2(one_before_split / length_all)
+    else:
+        npmi_matrix[0, 1] = -1.0
+
+    if zero_after_split > 0 and total_zeros > 0 and len_part2 > 0:
+        denominator = len_part2 * total_zeros
+        pmi_val = zero_after_split * length_all / denominator
+        pmi_matrix[1, 0] = (zero_after_split / length_all) * np.log2(pmi_val)
+        npmi_matrix[1, 0] = -np.log2(pmi_val) / np.log2(zero_after_split / length_all)
+    else:
+        npmi_matrix[1, 0] = -1.0
+
+    if one_after_split > 0 and total_ones > 0 and len_part2 > 0:
+        denominator = len_part2 * total_ones
+        pmi_val = one_after_split * length_all / denominator
+        pmi_matrix[1, 1] = (one_after_split / length_all) * np.log2(pmi_val)
+        npmi_matrix[1, 1] = -np.log2(pmi_val) / np.log2(one_after_split / length_all)
+    else:
+        npmi_matrix[1, 1] = -1.0
+
+    mi_value = pmi_matrix.sum()
+    normalized_pmi = (
+        npmi_matrix[0, 0] * zero_before_split
+        - npmi_matrix[0, 1] * one_before_split
+        - npmi_matrix[1, 0] * zero_after_split
+        + npmi_matrix[1, 1] * one_after_split
+    ) / length_all
+    return mi_value, normalized_pmi
+
+
+@numba.jit(nopython=True, cache=True)
+def maxpmi(ampbi, n):
+    """Find the split point that maximizes PMI.
+
+    Args:
+        ampbi: Binary sequence (0s and 1s)
+        n: If -1, search all transition points; otherwise use specified index
+
+    Returns:
+        (max_mi, max_normalized_pmi, split_indices)
+    """
+    zero_count = (ampbi == 0).sum()
+    one_count = (ampbi == 1).sum()
+
+    if zero_count == 0 or one_count == 0:
+        return -1.0, -1.0, np.array([-1], dtype=np.int64)
+
+    if n != -1:
+        mi_val, normalized_pmi_val = pmi(ampbi, n)
+        return mi_val, normalized_pmi_val, np.array([n], dtype=np.int64)
+
+    transition_indices = np.where(ampbi[1:] - ampbi[:-1] == 1)[0]
+    if len(transition_indices) == 0:
+        return -1.0, -1.0, np.array([-1], dtype=np.int64)
+
+    results = np.zeros((len(transition_indices), 2))
+    for i in range(len(transition_indices)):
+        results[i, 0], results[i, 1] = pmi(ampbi, transition_indices[i] + 1)
+
+    max_mi_value = np.max(results[:, 0])
+    best_indices = np.where(results[:, 0] == max_mi_value)[0]
+
+    final_split_points = transition_indices[best_indices] + 1
+    return max_mi_value, results[best_indices[0], 1], final_split_points
+
+
+@numba.jit(nopython=True, cache=True)
+def calculate_general(xsquare, n):
+    """Calculate general form for n-dimensional Gaussian integral.
+
+    Uses the formula for computing integrals of x^n * exp(-x^2/2*sigma^2).
+    """
+    if n == 3:
+        return 1.0 / xsquare
+    if n == 2:
+        return np.sqrt(np.pi / 2.0 / xsquare)
+    if n % 2 == 1:
+        m = (n - 3) / 2.0
+        prod = 1.0
+        for val in range(2, n - 1, 2):
+            prod *= val
+        return prod * (1.0 / xsquare) ** (m + 1.0)
+    else:
+        m = (n - 2) / 2.0
+        prod = 1.0
+        for val in range(1, n - 2, 2):
+            prod *= val
+        return prod * np.sqrt(np.pi / 2.0 / xsquare) * (1.0 / xsquare) ** m