asdasd

gistfile1.txt

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal

import numpy as np
import pandas as pd
from scipy.interpolate import griddata
from scipy.spatial.distance import pdist
from typing_extensions import Self

from hank.point import Point3d

if TYPE_CHECKING:
    pass


@dataclass
class Data:
    df: pd.DataFrame
    value_columns: tuple[str, ...]

    @classmethod
    def generate_scattered(
        cls,
        far_corner: Point3d,
        value: dict[str, tuple[float, float]],
        *,
        points: int = 500,
        rng: int | np.random.Generator | None = None,
    ) -> Self:
        assert not ("x" in value and "y" in value and "z" in value)
        rng = np.random.default_rng(rng)
        x, y, z = far_corner
        xs = np.random.uniform(0, x, size=points)
        ys = np.random.uniform(0, y, size=points)
        zs = np.random.uniform(0, z, size=points)
        values = {
            key: np.random.uniform(*value_range, size=points)
            for key, value_range in value.items()
        }
        return cls(
            df=pd.DataFrame({"x": xs, "y": ys, "z": zs, **values}),
            value_columns=tuple(values.keys()),
        )

    def as_grid(
        self,
        *,
        minimum: tuple[float, float, float] | None = None,
        maximum: tuple[float, float, float] | None = None,
        density: (
            tuple[Literal["points_per_axis"], int]
            | tuple[Literal["distance"], float]
            | tuple[Literal["percentile_min_distance"], float]
        ) = ("percentile_min_distance", 0.10),
    ) -> Data:
        if minimum is None:
            minimum = tuple(self.df[["x", "y", "z"]].min())
        if maximum is None:
            maximum = tuple(self.df[["x", "y", "z"]].max())

        match density:
            case ("points_per_axis", n_points):
                steps = (n_points, n_points, n_points)
            case ("distance", distance):
                steps = (
                    int(np.ceil((_max - _min) / distance)) + 1
                    for _min, _max in zip(minimum, maximum)
                )
            case ("percentile_min_distance", percentile):
                distance = np.percentile(pdist(self.df[["x", "y", "z"]]), percentile)
                steps = tuple(
                    int(np.ceil((_max - _min) / distance)) + 1
                    for _min, _max in zip(minimum, maximum)
                )
            case _:
                raise ValueError("Invalid density specification")

        xi, yi, zi = (
            np.linspace(_min, _max, _step)
            for _min, _max, _step in zip(minimum, maximum, steps)
        )
        grid_x, grid_y, grid_z = np.meshgrid(xi, yi, zi)

        grid_values = {}
        for value_column in self.value_columns:
            grid_values[value_column] = griddata(
                points=self.df[["x", "y", "z"]],
                values=self.df[value_column],
                xi=(grid_x, grid_y, grid_z),
                method="nearest",
            )

        print(grid_values)

        return self.__class__(
            df=pd.DataFrame(
                {
                    "x": grid_x.flatten(),
                    "y": grid_y.flatten(),
                    "z": grid_z.flatten(),
                    **{k: v.flatten() for k, v in grid_values.items()},
                }
            ),
            value_columns=self.value_columns,
        )

if __name__ == "__main__":
    import matplotlib.pyplot as plt

from hank import Data

data = Data.generate_scattered((1, 1, 1), value={"k": (0, 1), "v": (0, 1)}, points=5000)
uniform = data.as_grid(
    # density=("points_per_axis", 20)
)  # Use a smaller grid for faster plotting

fig = plt.figure(figsize=(12, 5))

# --- Plot 1: Original Scattered Data ---
ax1 = fig.add_subplot(121, projection="3d")
ax1.set_title("Original Scattered")
ax1.scatter(
    data.df["x"],
    data.df["y"],
    data.df["z"],
    c=data.df["k"],
    s=data.df["v"] * 20,
    alpha=0.6,
)

# --- Plot 2: Uniform Gridded Data ---
ax2 = fig.add_subplot(122, projection="3d")
ax2.set_title("Uniform Grid")
ax2.scatter(
    uniform.df["x"],
    uniform.df["y"],
    uniform.df["z"],
    c=uniform.df["k"],
    s=uniform.df["v"] * 20,
    alpha=0.6,
)

plt.tight_layout()
plt.show()