py4dgeo 1.0.0__cp314-cp314t-win_amd64.whl

py4dgeo/pbm3c2.py

@@ -0,0 +1,694 @@
+import numpy as np
+import pandas as pd
+from sklearn.ensemble import RandomForestClassifier
+from scipy.spatial import cKDTree
+from py4dgeo.util import Py4DGeoError
+from py4dgeo.epoch import Epoch
+import matplotlib.pyplot as plt
+from matplotlib.patches import FancyArrowPatch
+from mpl_toolkits.mplot3d import proj3d
+
+try:
+    from tqdm import tqdm
+except ImportError:
+
+    def tqdm(iterable, *args, **kwargs):
+        return iterable
+
+
+class PBM3C2:
+    """
+    Correspondence-driven plane-based M3C2 for lower uncertainty in 3D topographic change quantification.
+
+    This class implements the PBM3C2 algorithm as described in Zahs et al. (2022).
+    """
+
+    def __init__(self, registration_error=0.0):
+        self.clf = RandomForestClassifier(n_estimators=100, random_state=42)
+        self.registration_error = registration_error
+        self.epoch0_segment_metrics = None
+        self.epoch1_segment_metrics = None
+        self.epoch0_segments = None
+        self.epoch1_segments = None
+        self.correspondences = None
+        self.epoch0_id_mapping = None
+        self.epoch1_id_mapping = None
+        self.epoch0_reverse_mapping = None
+        self.epoch1_reverse_mapping = None
+
+    @staticmethod
+    def preprocess_epochs(epoch0, epoch1, correspondences_file):
+        """
+        Assign globally unique segment IDs using independent sequential numbering.
+        Map correspondence file IDs to the new ID scheme.
+
+        Parameters
+        ----------
+        epoch0, epoch1 : Epoch
+            Input epochs with segment_id in additional_dimensions
+        correspondences_file : str
+            Path to CSV file with correspondence data
+
+        Returns
+        -------
+        tuple
+            (processed_epoch0, processed_epoch1, remapped_correspondences,
+             epoch0_id_mapping, epoch1_id_mapping,
+             epoch0_reverse_mapping, epoch1_reverse_mapping)
+        """
+        print("Assigning globally unique segment IDs...")
+
+        orig_ids0 = np.unique(epoch0.additional_dimensions["segment_id"])
+        orig_ids1 = np.unique(epoch1.additional_dimensions["segment_id"])
+
+        print(
+            f"  Epoch 0: {len(orig_ids0)} unique segments (range: {orig_ids0.min()}-{orig_ids0.max()})"
+        )
+        print(
+            f"  Epoch 1: {len(orig_ids1)} unique segments (range: {orig_ids1.min()}-{orig_ids1.max()})"
+        )
+
+        # Load correspondence file
+        try:
+            correspondences_arr = np.genfromtxt(
+                correspondences_file, delimiter=",", dtype=np.float64
+            )
+            if correspondences_arr.ndim == 1:
+                correspondences_arr = correspondences_arr.reshape(1, -1)
+        except Exception as e:
+            raise Py4DGeoError(
+                f"Failed to read correspondence file '{correspondences_file}': {e}"
+            )
+
+        if correspondences_arr.shape[1] < 2:
+            raise Py4DGeoError(
+                f"The correspondence file must contain at least two columns (got {correspondences_arr.shape[1]})."
+            )
+
+        if not np.issubdtype(correspondences_arr.dtype, np.number):
+            raise Py4DGeoError(
+                "The correspondence file appears to contain non-numeric data."
+            )
+
+        corr_ids0 = correspondences_arr[:, 0]
+        corr_ids1 = correspondences_arr[:, 1]
+
+        # Validate correspondence IDs
+        invalid_ids0 = ~np.isin(corr_ids0, orig_ids0)
+        invalid_ids1 = ~np.isin(corr_ids1, orig_ids1)
+
+        if invalid_ids0.any():
+            invalid_list = corr_ids0[invalid_ids0]
+            print(
+                f"  Warning: {invalid_ids0.sum()} epoch0 IDs in correspondences don't exist in epoch0 segments"
+            )
+            print(f"    First 10 invalid IDs: {invalid_list[:10]}")
+
+        if invalid_ids1.any():
+            invalid_list = corr_ids1[invalid_ids1]
+            print(
+                f"  Warning: {invalid_ids1.sum()} epoch1 IDs in correspondences don't exist in epoch1 segments"
+            )
+            print(f"    First 10 invalid IDs: {invalid_list[:10]}")
+
+        # Filter invalid correspondences
+        valid_mask = ~invalid_ids0 & ~invalid_ids1
+        if not valid_mask.all():
+            print(f"  Filtering out {(~valid_mask).sum()} invalid correspondence pairs")
+            correspondences_arr = correspondences_arr[valid_mask]
+
+        if len(correspondences_arr) == 0:
+            raise Py4DGeoError(
+                "No valid correspondences remain after filtering. "
+                "Please verify that segment IDs in the correspondence file match those in the input epochs."
+            )
+
+        # Create new independent ID mappings
+        new_ids0 = np.arange(1, len(orig_ids0) + 1, dtype=np.int64)
+        epoch0_id_mapping = dict(zip(orig_ids0, new_ids0))
+        epoch0_reverse_mapping = dict(zip(new_ids0, orig_ids0))
+
+        start_id_epoch1 = len(orig_ids0) + 1
+        new_ids1 = np.arange(
+            start_id_epoch1, start_id_epoch1 + len(orig_ids1), dtype=np.int64
+        )
+        epoch1_id_mapping = dict(zip(orig_ids1, new_ids1))
+        epoch1_reverse_mapping = dict(zip(new_ids1, orig_ids1))
+
+        print(f"  Assigned new IDs for Epoch 0: 1 to {len(orig_ids0)}")
+        print(
+            f"  Assigned new IDs for Epoch 1: {start_id_epoch1} to {start_id_epoch1 + len(orig_ids1) - 1}"
+        )
+
+        # Apply new IDs to epoch0
+        old_seg_ids0 = epoch0.additional_dimensions["segment_id"]
+        new_seg_ids0 = np.array([epoch0_id_mapping[sid] for sid in old_seg_ids0])
+        new_add_dims0 = epoch0.additional_dimensions.copy()
+        new_add_dims0["segment_id"] = new_seg_ids0
+        new_epoch0 = Epoch(
+            cloud=epoch0.cloud.copy(), additional_dimensions=new_add_dims0
+        )
+
+        # Apply new IDs to epoch1
+        old_seg_ids1 = epoch1.additional_dimensions["segment_id"]
+        new_seg_ids1 = np.array([epoch1_id_mapping[sid] for sid in old_seg_ids1])
+        new_add_dims1 = epoch1.additional_dimensions.copy()
+        new_add_dims1["segment_id"] = new_seg_ids1
+        new_epoch1 = Epoch(
+            cloud=epoch1.cloud.copy(), additional_dimensions=new_add_dims1
+        )
+
+        # Remap correspondence IDs
+        remapped_corr = correspondences_arr.copy()
+        remapped_corr[:, 0] = np.array(
+            [epoch0_id_mapping[int(cid)] for cid in correspondences_arr[:, 0]]
+        )
+        remapped_corr[:, 1] = np.array(
+            [epoch1_id_mapping[int(cid)] for cid in correspondences_arr[:, 1]]
+        )
+
+        print(f"  Remapped {len(remapped_corr)} correspondence pairs to new ID scheme")
+        print("Preprocessing complete.\n")
+
+        return (
+            new_epoch0,
+            new_epoch1,
+            remapped_corr,
+            epoch0_id_mapping,
+            epoch1_id_mapping,
+            epoch0_reverse_mapping,
+            epoch1_reverse_mapping,
+        )
+
+    def _get_segments(self, epoch):
+        """Extract individual segments from an epoch."""
+        segment_id_array = epoch.additional_dimensions["segment_id"]
+        unique_segment_ids = np.unique(segment_id_array)
+
+        segments_dict = {}
+        for seg_id in unique_segment_ids:
+            indices = np.where(segment_id_array == seg_id)[0]
+            if len(indices) > 0:
+                segments_dict[seg_id] = {"points": epoch.cloud[indices].copy()}
+
+        return segments_dict
+
+    def _create_segment_metrics(self, segments):
+        """Calculate geometric metrics for each segment."""
+        metrics_list = []
+
+        for segment_id, data in tqdm(segments.items(), desc="Extracting Features"):
+            points = data["points"]
+            if points.shape[0] < 3:
+                continue
+
+            cog = np.mean(points, axis=0)
+            centered_points = points - cog
+            cov_matrix = np.cov(centered_points, rowvar=False)
+            eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
+
+            sort_indices = np.argsort(eigenvalues)[::-1]
+            eigenvalues = eigenvalues[sort_indices]
+            eigenvectors = eigenvectors[:, sort_indices]
+
+            normal = eigenvectors[:, 2]
+            distances_to_plane = np.dot(centered_points, normal)
+            roughness = np.std(distances_to_plane)
+
+            e1, e2, e3 = eigenvalues
+            sum_eigenvalues = e1 + e2 + e3
+            if sum_eigenvalues == 0:
+                linearity = planarity = sphericity = 0
+            else:
+                linearity = (e1 - e2) / sum_eigenvalues
+                planarity = (e2 - e3) / sum_eigenvalues
+                sphericity = e3 / sum_eigenvalues
+
+            metrics_list.append(
+                {
+                    "segment_id": segment_id,
+                    "cog_x": cog[0],
+                    "cog_y": cog[1],
+                    "cog_z": cog[2],
+                    "normal_x": normal[0],
+                    "normal_y": normal[1],
+                    "normal_z": normal[2],
+                    "linearity": linearity,
+                    "planarity": planarity,
+                    "sphericity": sphericity,
+                    "roughness": roughness,
+                    "num_points": points.shape[0],
+                }
+            )
+
+        return pd.DataFrame(metrics_list).set_index("segment_id")
+
+    def _create_feature_array(self, df_t1, df_t2, correspondences):
+        """Create feature vectors for segment pairs."""
+        features = []
+
+        if isinstance(correspondences, np.ndarray):
+            pairs = correspondences
+        else:
+            pairs = correspondences.values
+
+        for row in pairs:
+            id1, id2 = int(row[0]), int(row[1])
+
+            if id1 not in df_t1.index or id2 not in df_t2.index:
+                continue
+
+            metrics1 = df_t1.loc[id1]
+            metrics2 = df_t2.loc[id2]
+
+            cog1 = metrics1[["cog_x", "cog_y", "cog_z"]].values
+            cog2 = metrics2[["cog_x", "cog_y", "cog_z"]].values
+            cog_dist = np.linalg.norm(cog1 - cog2)
+
+            normal1 = metrics1[["normal_x", "normal_y", "normal_z"]].values
+            normal2 = metrics2[["normal_x", "normal_y", "normal_z"]].values
+            dot_product = np.clip(np.dot(normal1, normal2), -1.0, 1.0)
+            normal_angle = np.arccos(dot_product)
+
+            roughness_diff = abs(metrics1["roughness"] - metrics2["roughness"])
+
+            features.append([cog_dist, normal_angle, roughness_diff])
+
+        return np.array(features)
+
+    def train(self, correspondences):
+        """Train Random Forest classifier on labeled correspondences."""
+        positives = correspondences[correspondences[:, 2] == 1]
+        negatives = correspondences[correspondences[:, 2] == 0]
+
+        X_pos = self._create_feature_array(
+            self.epoch0_segment_metrics, self.epoch1_segment_metrics, positives
+        )
+        X_neg = self._create_feature_array(
+            self.epoch0_segment_metrics, self.epoch1_segment_metrics, negatives
+        )
+
+        if X_pos.shape[0] == 0 or X_neg.shape[0] == 0:
+            raise ValueError("Training data is missing positive or negative examples.")
+
+        X = np.vstack([X_pos, X_neg])
+        y = np.array([1] * len(X_pos) + [0] * len(X_neg))
+
+        self.clf.fit(X, y)
+
+    def apply(self, apply_ids, search_radius=3.0):
+        """Apply trained classifier to find correspondences."""
+        epoch1_cogs = self.epoch1_segment_metrics[["cog_x", "cog_y", "cog_z"]].values
+        kdtree = cKDTree(epoch1_cogs)
+        found_correspondences = []
+
+        for apply_id in tqdm(apply_ids, desc="Applying Classifier"):
+            if apply_id not in self.epoch0_segment_metrics.index:
+                continue
+
+            cog0 = self.epoch0_segment_metrics.loc[apply_id][
+                ["cog_x", "cog_y", "cog_z"]
+            ].values
+            indices = kdtree.query_ball_point(cog0, r=search_radius)
+
+            if len(indices) == 0:
+                continue
+
+            candidate_ids = self.epoch1_segment_metrics.index[indices]
+            apply_df = pd.DataFrame(
+                {"id1": [apply_id] * len(candidate_ids), "id2": candidate_ids}
+            )
+
+            X_apply = self._create_feature_array(
+                self.epoch0_segment_metrics, self.epoch1_segment_metrics, apply_df
+            )
+
+            if len(X_apply) == 0:
+                continue
+
+            probabilities = self.clf.predict_proba(X_apply)[:, 1]
+            best_match_idx = np.argmax(probabilities)
+            found_correspondences.append([apply_id, candidate_ids[best_match_idx]])
+
+        self.correspondences = pd.DataFrame(
+            found_correspondences, columns=["epoch0_segment_id", "epoch1_segment_id"]
+        )
+
+    def _calculate_cog_distance(self, segment1_id, segment2_id):
+        """Calculate CoG distance and level of detection."""
+        metrics1 = self.epoch0_segment_metrics.loc[segment1_id]
+        metrics2 = self.epoch1_segment_metrics.loc[segment2_id]
+
+        cog1 = metrics1[["cog_x", "cog_y", "cog_z"]].values
+        cog2 = metrics2[["cog_x", "cog_y", "cog_z"]].values
+        normal1 = metrics1[["normal_x", "normal_y", "normal_z"]].values
+
+        dist = np.dot(cog2 - cog1, normal1)
+
+        sigma1_sq = metrics1["roughness"] ** 2
+        sigma2_sq = metrics2["roughness"] ** 2
+        n1 = metrics1["num_points"]
+        n2 = metrics2["num_points"]
+
+        if n1 == 0 or n2 == 0:
+            lod = np.nan
+        else:
+            lod = (
+                1.96 * np.sqrt(sigma1_sq / n1 + sigma2_sq / n2)
+                + self.registration_error
+            )
+
+        return dist, lod
+
+    def run(self, epoch0, epoch1, correspondences_file, apply_ids, search_radius=3.0):
+        """
+        Execute complete PBM3C2 workflow.
+
+        Parameters
+        ----------
+        epoch0, epoch1 : Epoch
+            Input point cloud epochs with segment_id
+        correspondences_file : str
+            Path to CSV with training correspondences
+        apply_ids : array-like
+            Segment IDs to find correspondences for (using original IDs)
+        search_radius : float
+            Spatial search radius in meters
+
+        Returns
+        -------
+        DataFrame
+            Correspondences with distances and uncertainties (using original IDs)
+        """
+        print("=" * 60)
+        print("PBM3C2 Processing Pipeline")
+        print("=" * 60)
+
+        print("\n[1/6] Preprocessing epochs and correspondences...")
+        (
+            epoch0_processed,
+            epoch1_processed,
+            correspondences_for_training,
+            self.epoch0_id_mapping,
+            self.epoch1_id_mapping,
+            self.epoch0_reverse_mapping,
+            self.epoch1_reverse_mapping,
+        ) = self.preprocess_epochs(epoch0, epoch1, correspondences_file)
+
+        print("[2/6] Loading and processing segments...")
+        self.epoch0_segments = self._get_segments(epoch0_processed)
+        self.epoch1_segments = self._get_segments(epoch1_processed)
+        print(
+            f"  Loaded {len(self.epoch0_segments)} segments from epoch 0, {len(self.epoch1_segments)} from epoch 1"
+        )
+
+        print("\n[3/6] Extracting geometric features...")
+        self.epoch0_segment_metrics = self._create_segment_metrics(self.epoch0_segments)
+        self.epoch1_segment_metrics = self._create_segment_metrics(self.epoch1_segments)
+        print(
+            f"  Computed metrics for {len(self.epoch0_segment_metrics)} + {len(self.epoch1_segment_metrics)} segments"
+        )
+
+        print("\n[4/6] Training Random Forest classifier...")
+        self.train(correspondences_for_training)
+        print(f"  Classifier trained on {len(correspondences_for_training)} pairs")
+
+        print("\n[5/6] Finding correspondences...")
+        remapped_apply_ids = [
+            self.epoch0_id_mapping[orig_id]
+            for orig_id in apply_ids
+            if orig_id in self.epoch0_id_mapping
+        ]
+
+        if len(remapped_apply_ids) < len(apply_ids):
+            print(
+                f"  Warning: {len(apply_ids) - len(remapped_apply_ids)} apply_ids not found in epoch0"
+            )
+
+        self.apply(apply_ids=remapped_apply_ids, search_radius=search_radius)
+
+        print("\n[6/6] Calculating M3C2 distances and uncertainties...")
+        if self.correspondences is None or self.correspondences.empty:
+            print("  Warning: No correspondences were found.")
+            return self.correspondences
+
+        distances, uncertainties = [], []
+        for _, row in self.correspondences.iterrows():
+            id1, id2 = row["epoch0_segment_id"], row["epoch1_segment_id"]
+            dist, lod = self._calculate_cog_distance(id1, id2)
+            distances.append(dist)
+            uncertainties.append(lod)
+
+        self.correspondences["distance"] = distances
+        self.correspondences["uncertainty"] = uncertainties
+
+        print("\n[Final] Mapping results back to original segment IDs...")
+        self.correspondences["epoch0_original_id"] = self.correspondences[
+            "epoch0_segment_id"
+        ].map(self.epoch0_reverse_mapping)
+        self.correspondences["epoch1_original_id"] = self.correspondences[
+            "epoch1_segment_id"
+        ].map(self.epoch1_reverse_mapping)
+
+        cols = [
+            "epoch0_original_id",
+            "epoch1_original_id",
+            "epoch0_segment_id",
+            "epoch1_segment_id",
+            "distance",
+            "uncertainty",
+        ]
+        self.correspondences = self.correspondences[cols]
+
+        print("=" * 60)
+        print(f"Processing complete! Found {len(self.correspondences)} matches")
+        print("=" * 60)
+
+        return self.correspondences
+
+    def visualize_correspondences(
+        self,
+        epoch0_segment_id=None,
+        use_original_ids=True,
+        show_all=False,
+        num_samples=10,
+        figsize=(12, 10),
+        elev=30,
+        azim=45,
+    ):
+        """
+        Visualize matched plane segments and their correspondences.
+
+        Parameters
+        ----------
+        epoch0_segment_id : int, optional
+            Specific segment ID to visualize (original or new ID based on use_original_ids)
+        use_original_ids : bool, optional
+            If True, interpret epoch0_segment_id as original ID (default: True)
+        show_all : bool, optional
+            Plot all correspondences (default: False)
+        num_samples : int, optional
+            Number of random samples if show_all=False (default: 10)
+        figsize : tuple, optional
+            Figure size in inches
+        elev : float, optional
+            Elevation angle for 3D view in degrees
+        azim : float, optional
+            Azimuth angle for 3D view in degrees
+
+        Returns
+        -------
+        tuple
+            (fig, ax) matplotlib figure and axis objects
+        """
+        if self.correspondences is None or self.correspondences.empty:
+            raise ValueError("No correspondences found. Run run() first.")
+
+        n_corr = len(self.correspondences)
+        zoom_in = False
+
+        if epoch0_segment_id is not None:
+            if (
+                use_original_ids
+                and "epoch0_original_id" in self.correspondences.columns
+            ):
+                corr_sample = self.correspondences[
+                    self.correspondences["epoch0_original_id"] == epoch0_segment_id
+                ]
+                title_text = (
+                    f"PBM3C2 Correspondence for Original Segment ID {epoch0_segment_id}"
+                )
+            else:
+                corr_sample = self.correspondences[
+                    self.correspondences["epoch0_segment_id"] == epoch0_segment_id
+                ]
+                title_text = f"PBM3C2 Correspondence for Segment ID {epoch0_segment_id}"
+
+            if corr_sample.empty:
+                raise ValueError(
+                    f"Segment ID {epoch0_segment_id} not found in correspondences."
+                )
+            zoom_in = True
+
+        elif show_all:
+            corr_sample = self.correspondences
+            title_text = f"PBM3C2 Correspondences (showing all {n_corr})"
+
+        else:
+            if n_corr > num_samples:
+                sample_indices = np.random.choice(n_corr, num_samples, replace=False)
+                corr_sample = self.correspondences.iloc[sample_indices]
+                title_text = (
+                    f"PBM3C2 Correspondences (showing {num_samples} of {n_corr})"
+                )
+            else:
+                corr_sample = self.correspondences
+                title_text = f"PBM3C2 Correspondences (showing all {n_corr})"
+
+        fig = plt.figure(figsize=figsize)
+        ax = fig.add_subplot(111, projection="3d")
+
+        class Arrow3D(FancyArrowPatch):
+            def __init__(self, xs, ys, zs, *args, **kwargs):
+                super().__init__((0, 0), (0, 0), *args, **kwargs)
+                self._verts3d = xs, ys, zs
+
+            def do_3d_projection(self, renderer=None):
+                xs3d, ys3d, zs3d = self._verts3d
+                xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
+                self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
+                return np.min(zs)
+
+        epoch0_plotted = False
+        epoch1_plotted = False
+
+        for idx, row in corr_sample.iterrows():
+            id0 = row["epoch0_segment_id"]
+            id1 = row["epoch1_segment_id"]
+
+            if id0 not in self.epoch0_segments or id1 not in self.epoch1_segments:
+                continue
+
+            points0 = self.epoch0_segments[id0]["points"]
+            points1 = self.epoch1_segments[id1]["points"]
+
+            label0 = "Epoch 0" if not epoch0_plotted else ""
+            ax.scatter(
+                points0[:, 0],
+                points0[:, 1],
+                points0[:, 2],
+                c="blue",
+                s=5,
+                alpha=0.3,
+                label=label0,
+            )
+            if label0:
+                epoch0_plotted = True
+
+            label1 = "Epoch 1" if not epoch1_plotted else ""
+            ax.scatter(
+                points1[:, 0],
+                points1[:, 1],
+                points1[:, 2],
+                c="red",
+                s=5,
+                alpha=0.3,
+                label=label1,
+            )
+            if label1:
+                epoch1_plotted = True
+
+            cog0 = self.epoch0_segment_metrics.loc[id0][
+                ["cog_x", "cog_y", "cog_z"]
+            ].values
+            cog1 = self.epoch1_segment_metrics.loc[id1][
+                ["cog_x", "cog_y", "cog_z"]
+            ].values
+
+            ax.scatter(
+                *cog0, c="darkblue", s=100, marker="o", edgecolors="black", linewidths=2
+            )
+            ax.scatter(
+                *cog1, c="darkred", s=100, marker="o", edgecolors="black", linewidths=2
+            )
+
+            arrow = Arrow3D(
+                [cog0[0], cog1[0]],
+                [cog0[1], cog1[1]],
+                [cog0[2], cog1[2]],
+                mutation_scale=20,
+                lw=2,
+                arrowstyle="-|>",
+                color="green",
+                alpha=0.8,
+            )
+            ax.add_artist(arrow)
+
+            if "distance" in row:
+                mid_point = (cog0 + cog1) / 2
+                ax.text(
+                    mid_point[0],
+                    mid_point[1],
+                    mid_point[2],
+                    f'{row["distance"]:.3f}m',
+                    fontsize=8,
+                    color="green",
+                    weight="bold",
+                )
+
+        ax.set_xlabel("X [m]", fontsize=12)
+        ax.set_ylabel("Y [m]", fontsize=12)
+        ax.set_zlabel("Z [m]", fontsize=12)
+        ax.set_title(title_text, fontsize=14, weight="bold")
+
+        if epoch0_plotted or epoch1_plotted:
+            ax.legend(loc="upper right", fontsize=10)
+
+        ax.view_init(elev=elev, azim=azim)
+
+        if not zoom_in:
+            ax.set_box_aspect([1, 1, 1])
+
+        plt.tight_layout()
+
+        return fig, ax
+
+    def get_original_ids(self, new_epoch0_id=None, new_epoch1_id=None):
+        """Retrieve original segment IDs from new internal IDs."""
+        result = []
+
+        if new_epoch0_id is not None:
+            orig_id = self.epoch0_reverse_mapping.get(new_epoch0_id, None)
+            if orig_id is None:
+                print(f"Warning: New epoch0 ID {new_epoch0_id} not found in mapping")
+            result.append(orig_id)
+
+        if new_epoch1_id is not None:
+            orig_id = self.epoch1_reverse_mapping.get(new_epoch1_id, None)
+            if orig_id is None:
+                print(f"Warning: New epoch1 ID {new_epoch1_id} not found in mapping")
+            result.append(orig_id)
+
+        return tuple(result) if len(result) > 1 else (result[0] if result else None)
+
+    def get_new_ids(self, orig_epoch0_id=None, orig_epoch1_id=None):
+        """Retrieve new internal segment IDs from original IDs."""
+        result = []
+
+        if orig_epoch0_id is not None:
+            new_id = self.epoch0_id_mapping.get(orig_epoch0_id, None)
+            if new_id is None:
+                print(
+                    f"Warning: Original epoch0 ID {orig_epoch0_id} not found in mapping"
+                )
+            result.append(new_id)
+
+        if orig_epoch1_id is not None:
+            new_id = self.epoch1_id_mapping.get(orig_epoch1_id, None)
+            if new_id is None:
+                print(
+                    f"Warning: Original epoch1 ID {orig_epoch1_id} not found in mapping"
+                )
+            result.append(new_id)
+
+        return tuple(result) if len(result) > 1 else (result[0] if result else None)