GR-raytracing/src/ray.c

#include "ray.h"

#include "asymptotic.h"

#include <math.h>
#include <omp.h>
#include <stdlib.h>

int ray_pool_init(RayPool *p, size_t capacity) {
  if (p == NULL || capacity == 0)
    return -1;
  *p = (RayPool){.capacity = capacity};
#define RAY_ALLOC(field) (p->field = calloc(capacity, sizeof *p->field))
  if (!(RAY_ALLOC(t) && RAY_ALLOC(x0) && RAY_ALLOC(x1) && RAY_ALLOC(x2) &&
        RAY_ALLOC(p0) && RAY_ALLOC(p1) && RAY_ALLOC(p2) && RAY_ALLOC(observer) &&
        RAY_ALLOC(direction0) && RAY_ALLOC(direction1) && RAY_ALLOC(direction2) &&
        RAY_ALLOC(log_alpha_p0) && RAY_ALLOC(log_alpha_p0_0) &&
        RAY_ALLOC(activate_t) && RAY_ALLOC(steps) &&
        RAY_ALLOC(step_limit) && RAY_ALLOC(continuation) &&
        RAY_ALLOC(frame_id) && RAY_ALLOC(vertex_id) && RAY_ALLOC(status) &&
        RAY_ALLOC(endpoint))) {
    ray_pool_destroy(p);
    return -1;
  }
#undef RAY_ALLOC
  return 0;
}

int ray_pool_append(RayPool *p, const ObserverState *observer,
                    const double direction[3],
                    size_t frame_id, size_t vertex_id) {
  if (p == NULL || p->count == p->capacity)
    return -1;
  const size_t i = p->count;
  if (observer == NULL || direction == NULL)
    return -1;
  p->t[i] = observer->coordinate_time;
  p->activate_t[i] = observer->coordinate_time;
  p->observer[i] = observer;
  p->direction0[i] = direction[0];
  p->direction1[i] = direction[1];
  p->direction2[i] = direction[2];
  p->frame_id[i] = frame_id;
  p->vertex_id[i] = vertex_id;
  p->status[i] = RAY_POOL_PENDING;
  p->endpoint[i] = (RayEndpoint){.magnification = 1.0,
                                  .end_id = SPACETIME_END_NONE,
                                  .outcome = RAY_OUTCOME_INCOMPLETE,
                                  .reason = RAY_REASON_NONE,
                                  .stop_coordinate_time = NAN,
                                  .accepted_steps = 0,
                                  .final_x = {NAN, NAN, NAN},
                                  .final_Pi = {NAN, NAN, NAN},
                                  .final_log_alpha_p0 = NAN,
                                  .final_log_alpha_p0_0 = NAN,
                                  .threshold_value = NAN};
  p->step_limit[i] = 0;
  p->continuation[i] = 0;
  p->log_alpha_p0_0[i] = 0.0;
  ++p->count;
  return 0;
}

int ray_pool_append_continuation(RayPool *p, size_t frame_id,
                                 size_t vertex_id, double t, const double x[3],
                                 const double Pi[3], double log_alpha_p0,
                                 double log_alpha_p0_0, unsigned int steps,
                                 unsigned int limit) {
  if (p == NULL || p->count == p->capacity || x == NULL || Pi == NULL)
    return -1;
  const size_t i = p->count;
  p->t[i] = t;
  p->activate_t[i] = t;
  p->observer[i] = NULL;
  p->direction0[i] = p->direction1[i] = p->direction2[i] = 0.0;
  p->x0[i] = x[0]; p->x1[i] = x[1]; p->x2[i] = x[2];
  p->p0[i] = Pi[0]; p->p1[i] = Pi[1]; p->p2[i] = Pi[2];
  p->log_alpha_p0[i] = log_alpha_p0;
  p->log_alpha_p0_0[i] = log_alpha_p0_0;
  p->steps[i] = steps;
  p->step_limit[i] = limit;
  p->continuation[i] = 1;
  p->frame_id[i] = frame_id;
  p->vertex_id[i] = vertex_id;
  p->status[i] = RAY_POOL_PENDING;
  p->endpoint[i] = (RayEndpoint){.magnification = 1.0,
                                  .end_id = SPACETIME_END_NONE,
                                  .outcome = RAY_OUTCOME_INCOMPLETE,
                                  .reason = RAY_REASON_NONE,
                                  .stop_coordinate_time = NAN,
                                  .accepted_steps = steps,
                                  .final_x = {NAN, NAN, NAN},
                                  .final_Pi = {NAN, NAN, NAN},
                                  .final_log_alpha_p0 = NAN,
                                  .final_log_alpha_p0_0 = log_alpha_p0_0,
                                  .threshold_value = NAN};
  ++p->count;
  return 0;
}

void ray_pool_preroute(RayPool *p, const SpacetimeSource *source) {
  if (p == NULL || source == NULL)
    return;
#pragma omp parallel for schedule(static)
  for (size_t i = 0; i < p->count; ++i) {
    if (p->status[i] != RAY_POOL_PENDING || p->continuation[i])
      continue;
    AsymptoticRoute route;
    const AsymptoticStatus status = asymptotic_route_camera(
        source, p->observer[i],
        (double[]){p->direction0[i], p->direction1[i], p->direction2[i]},
        &route);
    if (status == ASYMPTOTIC_UNSUPPORTED || status == ASYMPTOTIC_INVALID) {
      p->endpoint[i].outcome = RAY_OUTCOME_INCOMPLETE;
      p->endpoint[i].reason = status == ASYMPTOTIC_UNSUPPORTED
                                  ? RAY_REASON_UNSUPPORTED
                                  : RAY_REASON_PROTOCOL_ERROR;
      p->endpoint[i].end_id = route.end_id;
      p->status[i] = RAY_POOL_FAILED;
      continue;
    }
    if (status == ASYMPTOTIC_TIME_RANGE_EXHAUSTED) {
      p->endpoint[i].outcome = RAY_OUTCOME_INCOMPLETE;
      p->endpoint[i].reason = RAY_REASON_TIME_RANGE_EXHAUSTED;
      p->endpoint[i].end_id = route.end_id;
      p->status[i] = RAY_POOL_TERMINATED;
      continue;
    }
    if (route.kind == ASYMPTOTIC_ROUTE_ESCAPED) {
      for (int axis = 0; axis < 3; ++axis)
        p->endpoint[i].n_infinity[axis] = route.n_infinity[axis];
      p->endpoint[i].frequency_ratio = route.frequency_ratio;
      p->endpoint[i].end_id = route.end_id;
      p->endpoint[i].outcome = RAY_OUTCOME_ESCAPED;
      p->endpoint[i].reason = RAY_REASON_NONE;
      p->status[i] = RAY_POOL_TERMINATED;
      continue;
    }
    if (route.kind == ASYMPTOTIC_ROUTE_TIME_RANGE_EXHAUSTED) {
      p->endpoint[i].outcome = RAY_OUTCOME_INCOMPLETE;
      p->endpoint[i].reason = RAY_REASON_TIME_RANGE_EXHAUSTED;
      p->endpoint[i].end_id = route.end_id;
      p->status[i] = RAY_POOL_TERMINATED;
      continue;
    }
    if (route.kind != ASYMPTOTIC_ROUTE_INSIDE &&
        route.kind != ASYMPTOTIC_ROUTE_ENTRY) {
      p->endpoint[i].outcome = RAY_OUTCOME_INCOMPLETE;
      p->endpoint[i].reason = RAY_REASON_PROTOCOL_ERROR;
      p->status[i] = RAY_POOL_FAILED;
      continue;
    }
    p->activate_t[i] = route.activate_t;
    p->x0[i] = route.x[0];
    p->x1[i] = route.x[1];
    p->x2[i] = route.x[2];
    p->p0[i] = route.Pi[0];
    p->p1[i] = route.Pi[1];
    p->p2[i] = route.Pi[2];
    p->log_alpha_p0[i] = route.log_alpha_p0;
    /* Camera-event reference, distinct from the entry-state L. */
    p->log_alpha_p0_0[i] = route.log_alpha_p0_camera;
  }
}

void ray_pool_activate_in_time_range(RayPool *p, const MetricSlab *slab) {
  for (size_t i = 0; i < p->count; ++i) {
    if (p->status[i] != RAY_POOL_PENDING || p->activate_t[i] > slab->t_hi ||
        p->activate_t[i] <= slab->t_lo)
      continue;
    p->t[i] = p->activate_t[i];
    /* Continuation rays keep the accepted-step count they already consumed. */
    if (!p->continuation[i])
      p->steps[i] = 0;
    p->status[i] = RAY_POOL_ACTIVE;
  }
}

void ray_pool_advance_active(RayPool *p, const MetricSlab *slab,
                             const GeodesicTraceConfig *config) {
  /* Samples are grouped by frame, so a time slab activates only part of the
   * pool. Give workers bounded chunks instead of one contiguous frame range;
   * dynamic assignment also balances varying strong-field integration costs.
   * Pool slots stay stable for endpoint installation after the slab sweep. */
#pragma omp parallel for schedule(dynamic, 32)
  for (size_t i = 0; i < p->count; ++i) {
    if (p->status[i] != RAY_POOL_ACTIVE)
      continue;
    GeodesicRayState s = {.coordinate_time = p->t[i],
                           .x = {p->x0[i], p->x1[i], p->x2[i]},
                           .Pi = {p->p0[i], p->p1[i], p->p2[i]},
                           .log_alpha_p0 = p->log_alpha_p0[i],
                           .log_alpha_p0_0 = p->log_alpha_p0_0[i],
                           .steps = p->steps[i]};
    GeodesicTraceConfig per_ray = *config;
    if (p->step_limit[i] != 0)
      per_ray.max_steps = p->step_limit[i];
    const GeodesicAdvanceResult result =
        geodesic_advance_past_ray(slab, &s, slab->t_lo, &per_ray, &p->endpoint[i]);
    p->t[i] = s.coordinate_time;
    p->x0[i] = s.x[0]; p->x1[i] = s.x[1]; p->x2[i] = s.x[2];
    p->p0[i] = s.Pi[0]; p->p1[i] = s.Pi[1]; p->p2[i] = s.Pi[2];
    p->log_alpha_p0[i] = s.log_alpha_p0;
    p->steps[i] = s.steps;
    if (result == GEODESIC_ADVANCE_TERMINATED)
      p->status[i] = p->endpoint[i].outcome == RAY_OUTCOME_UNRESOLVED
                         ? RAY_POOL_UNRESOLVED
                         : RAY_POOL_TERMINATED;
    else if (result == GEODESIC_ADVANCE_FAILED)
      p->status[i] = RAY_POOL_FAILED;
  }
}

int ray_pool_has_live(const RayPool *p) {
  if (p == NULL)
    return 0;
  for (size_t i = 0; i < p->count; ++i)
    if (p->status[i] == RAY_POOL_PENDING || p->status[i] == RAY_POOL_ACTIVE)
      return 1;
  return 0;
}

void ray_pool_destroy(RayPool *p) {
  if (p == NULL)
    return;
  free(p->t); free(p->x0); free(p->x1); free(p->x2); free(p->observer);
  free(p->direction0); free(p->direction1); free(p->direction2);
  free(p->p0); free(p->p1); free(p->p2); free(p->log_alpha_p0);
  free(p->log_alpha_p0_0);
  free(p->activate_t); free(p->steps); free(p->step_limit);
  free(p->continuation); free(p->frame_id); free(p->vertex_id);
  free(p->status);
  free(p->endpoint);
  *p = (RayPool){0};
}