machine learning numpy/pytorch utils for circular or radial variables

radial.py

from torch import nn
import torch
import torch.nn.functional as F
import numpy as np



def radneg2radpos(f):
    """convert radians in range [-pi,pi] to [0,2*pi]"""
    return np.where(f < 0, f + np.pi * 2, f)


def radpos2radneg(f):
    """convert radians in range [0,2*pi] to [-pi,pi]"""
    return np.where(f > np.pi, f - np.pi * 2, f)


def wrap_deg_neg(d):
    """Wrap large degrees to [-180, 180]."""
    d = np.where(d < -180, d + 360, d)
    d = np.where(d > 180, d - 360, d)
    return d
  
 def wrap_deg_pos(d):
  """Wrap large degrees to [0, 360]."""
    d = np.where(d < 0, d + 360, d)
    d = np.where(d > 360, d - 360, d)
    return d

def wrap_deg_neg_df(df):
    values = wrap_deg_neg(df)
    if isinstance(df, pd.DataFrame):
        return pd.DataFrame(values, index=df.index, columns=df.columns)
    elif isinstance(df, pd.Series):
        return pd.Series(values, index=df.index)
    else:
        return values



def circmean_deg(x):
    """circmean for degrees."""
    return np.rad2deg(scipy.stats.circmean(np.deg2rad(x)))



def softnorm(x: torch.Tensor) -> torch.Tensor:
    """
    Makes the vector components have a magnitude of 1.
    Similar to softmax which makes sure final channels adds to one. 
    Inputs:
    - x: tensor with shape [Batch, ..., components]
    """
    magnitude = torch.norm(x, dim=-1, keepdim=True)
    return x / magnitude