# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../notebooks/utils/02_ops.ipynb.

# %% auto 0
__all__ = ['vector_project', 'vector_project_pl']

# %% ../../../notebooks/utils/02_ops.ipynb 0
import xarray as xr
import polars as pl
from xarray_einstats import linalg
from .basic import df2ts

# %% ../../../notebooks/utils/02_ops.ipynb 1
def vector_project(v1: xr.DataArray, v2: xr.DataArray, dim="v_dim"):
    return xr.dot(v1, v2, dim=dim) / linalg.norm(v2, dims=dim)


def vector_project_pl(df: pl.DataFrame, v1_cols, v2_cols, name=None):
    v1 = df2ts(df, v1_cols).assign_coords(v_dim=["x", "y", "z"])
    v2 = df2ts(df, v2_cols).assign_coords(v_dim=["x", "y", "z"])
    result = vector_project(v1, v2, dim="v_dim")

    return df.with_columns(pl.Series(result.data).alias(name or "v_proj"))
