cepheid

defmodule Cepheid.Periodicity.Pdm do
  @moduledoc """
  Functions for Phase Disperion Minimization (PDM) analysis.
  """

  import Nx.Defn

  defnp sum_of_distances(phases, samples) do
    # phases and samples should be sorted in phase order

    # create pairs of adjacent points by slicing tensors
    phase_diff =
      Nx.slice_along_axis(phases, 1, Nx.size(phases) - 1) -
        Nx.slice_along_axis(phases, 0, Nx.size(phases) - 1)

    sample_diff =
      Nx.slice_along_axis(samples, 1, Nx.size(samples) - 1) -
        Nx.slice_along_axis(samples, 0, Nx.size(samples) - 1)

    # calculate euclidean distances
    distances = Nx.sqrt(Nx.pow(phase_diff, 2) + Nx.pow(sample_diff, 2))

    # return the sum of all distances - "string length"
    Nx.sum(distances)
  end

  defnp compute_phases(period, times_since_epoch) do
    # Calculate the phase for each time value
    Nx.remainder(times_since_epoch / period, 1)
  end

  @doc """
  Period-finding via the string length method using trial periods

  ## Parameters
  - `min_period`: Minimum period to test (units should be 1/time units of times_since_epoch).
  - `max_period`: Maximum period to test.
  - `step_size`: Step size for period testing.
  - `times_since_epoch`: A 1d tensor of time values of time units since epoch
  - `samples`: A 1d tensor of sample values.

  ## Returns
  - `estimated_period`: The period that minimizes the sum of distances.

  See https://ui.adsabs.harvard.edu/abs/1983MNRAS.203..917D/abstract for background
  """
  def pdm_string(min_period, max_period, step_size, times_since_epoch, samples) do
    trial_periods =
      Nx.linspace(min_period, max_period, n: round((max_period - min_period) / step_size) + 1)

    # get phase values for times given each trial period, saving the one that minimizes the sum of distances
    {estimated_period, _min_distance} =
      Enum.reduce(trial_periods, {nil, :infinity}, fn period, {best_period, best_distance} ->
        phases = compute_phases(period, times_since_epoch)

        # sort phases and samples
        indices = Nx.argsort(phases)
        sorted_phases = Nx.take(phases, indices)
        sorted_samples = Nx.take(samples, indices)

        distance = sum_of_distances(sorted_phases, sorted_samples)

        if distance < best_distance do
          {period, distance}
        else
          {best_period, best_distance}
        end
      end)

    # Return the period that minimizes the sum of distances
    estimated_period
  end
end