path: root/td_constraints.c

/*
 * Copyright 2017 Anton Khirnov <anton@khirnov.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <errno.h>
#include <math.h>

#include "common.h"
#include "pssolve.h"
#include "td_constraints.h"

struct TDConstraintEvalPriv {
    double *k_rtheta;
    double *dk_rtheta[2];
};

typedef void (*ConstraintEvalCB)(TDConstraintEvalContext *ctx,
                                 const double * const vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                                 double *out);

typedef struct TDFamilyDef {
    double (*eval_krt)   (TDConstraintEvalContext *ctx, double r, double theta);
    double (*eval_krt_dr)(TDConstraintEvalContext *ctx, double r, double theta);
    double (*eval_krt_dt)(TDConstraintEvalContext *ctx, double r, double theta);
    double a_converge;
    double a_diverge;

    const ConstraintEvalCB *constraint_eval;
} TDFamilyDef;

static void
constraint_eval_ham_confflat(TDConstraintEvalContext *ctx,
                             const double * const vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                             double *out)
{
    TDConstraintEvalPriv *priv = ctx->priv;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];
        const double tt = tan(theta);

        const double tt_normalized = tt / M_PI;
        const int    is_axis       = ABS(tt_normalized - round(tt_normalized)) < EPS;
        const double dtt_t         = ((int)tt_normalized) & 1 ? -1.0 : 1.0;

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            const double K_rt = priv->k_rtheta[idx];

            const double K_rr     = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_00][idx];
            const double K_pp     = vars[TD_CONSTRAINT_VAR_CURV_PHIPHI][PSSOLVE_DIFF_ORDER_00][idx];
            const double psi      = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_00][idx] + 1.0;
            const double dpsi_r   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_10][idx];
            const double dpsi_t   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_01][idx];
            const double d2psi_rr = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_20][idx];
            const double d2psi_tt = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_02][idx];

            const double Km[3][3] = {{ K_rr,                    K_rt,    0 },
                                     { K_rt / SQR(r), -(K_rr + K_pp),    0 },
                                     { 0,                          0, K_pp }};

            const double r2   = SQR(r);
            const double psi2 = SQR(psi);
            const double psi4 = SQR(psi2);

            // term that is singular on the axis, use l'Hospital rule
            const double dpsi_term = is_axis ? d2psi_tt / dtt_t : dpsi_t / tt;
            const double Rscal = -8.0 * (d2psi_rr + 2.0 * dpsi_r / r + d2psi_tt / r2 + dpsi_term / r2) / (psi4 * psi);

            double K2 = 0.0;
            for (int i = 0; i < 3; i++)
                for (int j = 0; j < 3; j++)
                    K2 += Km[i][j] * Km[j][i];

            out[idx] = Rscal - K2;
        }
    }
}

static void
constraint_eval_mom_r_confflat(TDConstraintEvalContext *ctx,
                               const double * const vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                               double *out)
{
    TDConstraintEvalPriv *priv = ctx->priv;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];
        const double tt = tan(theta);

        const double tt_normalized = tt / M_PI;
        const int    is_axis       = ABS(tt_normalized - round(tt_normalized)) < EPS;
        const double dtt_t         = ((int)tt_normalized) & 1 ? -1.0 : 1.0;

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            const double K_rt     = priv->k_rtheta[idx];
            const double dK_rt_t  = priv->dk_rtheta[1][idx];

            const double K_rr     = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_00][idx];
            const double dK_rr_r  = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_10][idx];
            const double psi      = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_00][idx] + 1.0;
            const double dpsi_r   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_10][idx];
            const double dpsi_t   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_01][idx];

            const double r2   = SQR(r);

            // term that is singular on the axis, use l'Hospital
            const double Krt_singular = is_axis ? dK_rt_t / dtt_t : K_rt / tt;

            out[idx] = 6.0 * (K_rr * dpsi_r + K_rt * dpsi_t / r2) / psi +
                       dK_rr_r + 3.0 * K_rr / r + Krt_singular / r2 + dK_rt_t / r2;
        }
    }
}

static void
constraint_eval_mom_t_confflat(TDConstraintEvalContext *ctx,
                               const double * const vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                               double *out)
{
    TDConstraintEvalPriv *priv = ctx->priv;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];
        const double tt = tan(theta);

        const double tt_normalized = tt / M_PI;
        const int    is_axis       = ABS(tt_normalized - round(tt_normalized)) < EPS;
        const double dtt_t         = ((int)tt_normalized) & 1 ? -1.0 : 1.0;

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            const double K_rt     = priv->k_rtheta[idx];
            const double dK_rt_r  = priv->dk_rtheta[0][idx];

            const double K_rr     = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_00][idx];
            const double dK_rr_t  = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_01][idx];
            const double K_pp     = vars[TD_CONSTRAINT_VAR_CURV_PHIPHI][PSSOLVE_DIFF_ORDER_00][idx];
            const double dK_pp_t  = vars[TD_CONSTRAINT_VAR_CURV_PHIPHI][PSSOLVE_DIFF_ORDER_01][idx];
            const double psi      = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_00][idx] + 1.0;
            const double dpsi_r   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_10][idx];
            const double dpsi_t   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_01][idx];

            // term that is singular on the axis, use l'Hospital
            const double K_singular = is_axis ? (dK_rr_t + 2.0 * dK_pp_t) / dtt_t : (K_rr + 2.0 * K_pp) / tt;

            out[idx] = -K_singular - 6.0 * (K_rr + K_pp) * dpsi_t / psi +
                       6.0 * K_rt * dpsi_r / psi - dK_pp_t - dK_rr_t + dK_rt_r + 2.0 * K_rt / r;
        }
    }
}

static const ConstraintEvalCB constraint_funcs_confflat[TD_CONSTRAINT_EQ_NB] = {
    [TD_CONSTRAINT_EQ_HAM]   = constraint_eval_ham_confflat,
    [TD_CONSTRAINT_EQ_MOM_0] = constraint_eval_mom_r_confflat,
    [TD_CONSTRAINT_EQ_MOM_1] = constraint_eval_mom_t_confflat,
};

static double k_rtheta_eval_time_antisym_cubic(TDConstraintEvalContext *ctx,
                                               double r, double theta)
{
    const double r2 = SQR(r);
    const double r3 = r * r2;
    return (r3 - r) * exp(-r2) * sin(2.0 * theta) * ctx->amplitude;
}

static double k_rtheta_eval_dr_time_antisym_cubic(TDConstraintEvalContext *ctx,
                                               double r, double theta)
{
    const double r2 = SQR(r);
    const double r3 = r * r2;
    return (3.0 * r2 - 1.0 - 2.0 * r  * (r3 - r)) * exp(-r2) * sin(2.0 * theta) * ctx->amplitude;
}

static double k_rtheta_eval_dt_time_antisym_cubic(TDConstraintEvalContext *ctx,
                                                  double r, double theta)
{
    const double r2 = SQR(r);
    const double r3 = r * r2;
    return 2.0 * (r3 - r) * exp(-r2) * cos(2.0 * theta) * ctx->amplitude;
}

static const TDFamilyDef time_antisym_cubic = {
    .eval_krt    = k_rtheta_eval_time_antisym_cubic,
    .eval_krt_dr = k_rtheta_eval_dr_time_antisym_cubic,
    .eval_krt_dt = k_rtheta_eval_dt_time_antisym_cubic,
    //.a_converge =  1.362473539196777,
    .a_converge =  1.362473,
    .a_diverge  =  1.362473545703125,
    .constraint_eval = constraint_funcs_confflat,
};

static double k_rtheta_eval_time_antisym_linear(TDConstraintEvalContext *ctx,
                                                double r, double theta)
{
    const double r2 = SQR(r);
    return ctx->amplitude * r * exp(-r2) * sin(2.0 * theta);
}

static double k_rtheta_eval_dr_time_antisym_linear(TDConstraintEvalContext *ctx,
                                                   double r, double theta)
{
    const double r2 = SQR(r);
    return ctx->amplitude * (1.0 - 2.0 * r2) * exp(-r2) * sin(2.0 * theta);
}

static double k_rtheta_eval_dt_time_antisym_linear(TDConstraintEvalContext *ctx,
                                                   double r, double theta)
{
    const double r2 = SQR(r);
    return 2.0 * ctx->amplitude * r * exp(-r2) * cos(2.0 * theta);
}

static const TDFamilyDef time_antisym_linear = {
    .eval_krt    = k_rtheta_eval_time_antisym_linear,
    .eval_krt_dr = k_rtheta_eval_dr_time_antisym_linear,
    .eval_krt_dt = k_rtheta_eval_dt_time_antisym_linear,
    //.a_converge  = 1.018918628476058,
    .a_converge  = 1.0189186,
    .a_diverge   = 1.018918628476968,
    .constraint_eval = constraint_funcs_confflat,
};

static const TDFamilyDef *td_families[] = {
    [TD_FAMILY_TIME_ANTISYM_CUBIC] = &time_antisym_cubic,
    [TD_FAMILY_TIME_ANTISYM_LINEAR] = &time_antisym_linear,
};

double tdi_constraint_eval_k_rtheta(TDConstraintEvalContext *ctx, double r, double theta)
{
    const TDFamilyDef *fd = td_families[ctx->family];
    return fd->eval_krt(ctx, r, theta);
}

double tdi_constraint_eval_dk_rtheta_r(TDConstraintEvalContext *ctx, double r, double theta)
{
    const TDFamilyDef *fd = td_families[ctx->family];
    return fd->eval_krt_dr(ctx, r, theta);
}

double tdi_constraint_eval_dk_rtheta_t(TDConstraintEvalContext *ctx, double r, double theta)
{
    const TDFamilyDef *fd = td_families[ctx->family];
    return fd->eval_krt_dt(ctx, r, theta);
}

int tdi_constraint_eval_init(TDConstraintEvalContext *ctx)
{
    TDConstraintEvalPriv *priv = ctx->priv;
    const TDFamilyDef *fd = td_families[ctx->family];
    int ret;

    ret = posix_memalign((void**)&priv->k_rtheta, 32,
                         sizeof(*priv->k_rtheta) * ctx->nb_coords[0] * ctx->nb_coords[1]);
    if (ret)
        return -ENOMEM;

    for (int i = 0; i < ARRAY_ELEMS(priv->dk_rtheta); i++) {
        ret = posix_memalign((void**)&priv->dk_rtheta[i], 32,
                             sizeof(*priv->dk_rtheta[i]) * ctx->nb_coords[0] * ctx->nb_coords[1]);
        if (ret)
            return -ENOMEM;
    }

    ctx->a_converge = fd->a_converge;
    ctx->a_diverge  = fd->a_diverge;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            priv->k_rtheta[idx]     = fd->eval_krt   (ctx, r, theta);
            priv->dk_rtheta[0][idx] = fd->eval_krt_dr(ctx, r, theta);
            priv->dk_rtheta[1][idx] = fd->eval_krt_dt(ctx, r, theta);
        }
    }

    return 0;
}

int tdi_constraint_eval_alloc(TDConstraintEvalContext **pctx, enum TDFamily family)
{
    TDConstraintEvalContext *ctx;
    int ret;

    if (family < 0 || family >= ARRAY_ELEMS(td_families))
        return -EINVAL;

    ctx = calloc(1, sizeof(*ctx));
    if (!ctx)
        return -ENOMEM;

    ctx->priv = calloc(1, sizeof(*ctx->priv));
    if (!ctx->priv) {
        tdi_constraint_eval_free(&ctx);
        return -ENOMEM;
    }

    ctx->family = family;

    *pctx = ctx;
    return 0;
}

void tdi_constraint_eval_free(TDConstraintEvalContext **pctx)
{
    TDConstraintEvalContext *ctx = *pctx;

    if (!ctx)
        return;

    if (ctx->priv) {
        free(ctx->priv->k_rtheta);
        free(ctx->priv->dk_rtheta[0]);
        free(ctx->priv->dk_rtheta[1]);
    }
    free(ctx->priv);

    free(ctx);
    *pctx = NULL;
}

int tdi_constraint_eval(TDConstraintEvalContext *ctx, enum TDConstraintEq eq,
                        const double * const vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                        double *out)
{
    const TDFamilyDef *fd = td_families[ctx->family];
    fd->constraint_eval[eq](ctx, vars, out);
    return 0;
}