Skip to content

Refactor: CubicSpline class stores derivatives instead of polynomial coefficients #2772

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jinzx10 merged 3 commits into from
Jul 31, 2023
Merged

Refactor: CubicSpline class stores derivatives instead of polynomial coefficients #2772

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
83889a0
CubicSpline add support for derivative
jinzx10 Jul 30, 2023
70026e5
function name fix
jinzx10 Jul 30, 2023
1a81602
Merge branch 'develop' into better_spline
jinzx10 Jul 30, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 61 additions & 96 deletions source/module_base/cubic_spline.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
#include "cubic_spline.h"

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstring>
#include <functional>

namespace ModuleBase
{
Expand All @@ -15,16 +17,12 @@ CubicSpline::~CubicSpline()
void CubicSpline::cleanup()
{
delete[] x_;
delete[] c0_;
delete[] c1_;
delete[] c2_;
delete[] c3_;
delete[] y_;
delete[] s_;

x_ = nullptr;
c0_ = nullptr;
c1_ = nullptr;
c2_ = nullptr;
c3_ = nullptr;
y_ = nullptr;
s_ = nullptr;
}

void CubicSpline::sanity_check(const int n,
Expand Down Expand Up @@ -65,35 +63,26 @@ void CubicSpline::build(const int n,

n_ = n;
x_ = new double[n];
y_ = new double[n];
std::memcpy(x_, x, sizeof(double) * n);
std::memcpy(y_, y, sizeof(double) * n);

c0_ = new double[n_ - 1];
c1_ = new double[n_ - 1];
c2_ = new double[n_ - 1];
c3_ = new double[n_ - 1];
// to be computed
s_ = new double[n];

if (n == 2 && bc_start == BoundaryCondition::periodic)
{
// in this case the polynomial is a constant
c0_[0] = y[0];
c1_[0] = c2_[0] = c3_[0] = 0.0;
{ // in this case the polynomial is a constant
s_[0] = s_[1] = 0.0;
}
else if (n == 3 && bc_start == BoundaryCondition::not_a_knot && bc_end == BoundaryCondition::not_a_knot)
{
// in this case two conditions coincide; simply build a parabola that passes through the three data points
{ // in this case two conditions coincide; simply build a parabola that passes through the three data points
double idx10 = 1. / (x[1] - x[0]);
double idx21 = 1. / (x[2] - x[1]);
double idx20 = 1. / (x[2] - x[0]);

c0_[0] = y[0];
c1_[0] = -y[0] * (idx10 + idx20) + y[1] * (idx21 + idx10) + y[2] * (idx20 - idx21);
c2_[0] = y[0] * idx10 * idx20 - y[1] * idx10 * idx21 + y[2] * idx20 * idx21;
c3_[0] = 0.0;

c0_[1] = y[1];
c1_[1] = -y[1] * (-idx10 + idx21) + y[0] * (idx20 - idx10) + y[2] * (idx21 - idx20);
c2_[1] = -y[1] * idx10 * idx21 + y[0] * idx10 * idx20 + y[2] * idx21 * idx20;
c3_[1] = 0.0;
s_[0] = -y[0] * (idx10 + idx20) + y[1] * (idx21 + idx10) + y[2] * (idx20 - idx21);
s_[1] = -y[1] * (-idx10 + idx21) + y[0] * (idx20 - idx10) + y[2] * (idx21 - idx20);
s_[2] = s_[1] + 2.0 * (-y[1] * idx10 + y[2] * idx20) + 2.0 * y[0] * idx10 * idx20 * (x[2] - x[1]);
}
else
{
Expand All @@ -104,7 +93,6 @@ void CubicSpline::build(const int n,
double* diag = new double[n];
double* subdiag = new double[n - 1];
double* supdiag = new double[n - 1];
double* s = new double[n];

is_uniform_ = true;
double dx_avg = (x[n - 1] - x[0]) / (n - 1);
Expand All @@ -126,7 +114,7 @@ void CubicSpline::build(const int n,
diag[i] = 2.0 * (dx[i - 1] + dx[i]);
supdiag[i] = dx[i - 1];
subdiag[i - 1] = dx[i];
s[i] = 3.0 * (dd[i - 1] * dx[i] + dd[i] * dx[i - 1]);
s_[i] = 3.0 * (dd[i - 1] * dx[i] + dd[i] * dx[i - 1]);
}

if (bc_start == BoundaryCondition::periodic)
Expand All @@ -136,125 +124,102 @@ void CubicSpline::build(const int n,
diag[0] = 2.0 * (dx[n - 2] + dx[0]);
supdiag[0] = dx[n - 2];
subdiag[n - 2] = dx[0];
s[0] = 3.0 * (dd[0] * dx[n - 2] + dd[n - 2] * dx[0]);
s_[0] = 3.0 * (dd[0] * dx[n - 2] + dd[n - 2] * dx[0]);
;
solve_cyctri(n - 1, diag, supdiag, subdiag, s);
s[n - 1] = s[0];
solve_cyctri(n - 1, diag, supdiag, subdiag, s_);
s_[n - 1] = s_[0];
}
else
{

switch (bc_start)
{
case BoundaryCondition::not_a_knot:
diag[0] = dx[1];
supdiag[0] = x[2] - x[0];
s[0] = (dd[0] * dx[1] * (dx[0] + 2 * (x[2] - x[0])) + dd[1] * dx[0] * dx[0]) / (x[2] - x[0]);
break;
case BoundaryCondition::first_deriv:
diag[0] = 1.0 * dx[0];
supdiag[0] = 0.0;
s[0] = deriv_start * dx[0];
s_[0] = deriv_start * dx[0];
break;
case BoundaryCondition::second_deriv:
diag[0] = 2.0 * dx[0];
supdiag[0] = 1.0 * dx[0];
s[0] = (3.0 * dd[0] - 0.5 * deriv_start * dx[0]) * dx[0];
s_[0] = (3.0 * dd[0] - 0.5 * deriv_start * dx[0]) * dx[0];
break;
default:
assert(false && "CubicSpline: BoundaryCondition error!"); // not supposed to happend
default: // BoundaryCondition::not_a_knot
diag[0] = dx[1];
supdiag[0] = x[2] - x[0];
s_[0] = (dd[0] * dx[1] * (dx[0] + 2 * (x[2] - x[0])) + dd[1] * dx[0] * dx[0]) / (x[2] - x[0]);
}

switch (bc_end)
{
case BoundaryCondition::not_a_knot:
diag[n - 1] = dx[n - 3];
subdiag[n - 2] = x[n - 1] - x[n - 3];
s[n - 1] = (dd[n - 2] * dx[n - 3] * (dx[n - 2] + 2 * (x[n - 1] - x[n - 3]))
+ dd[n - 3] * dx[n - 2] * dx[n - 2])
/ (x[n - 1] - x[n - 3]);
break;
case BoundaryCondition::first_deriv:
diag[n - 1] = 1.0 * dx[n - 2];
subdiag[n - 2] = 0.0;
s[n - 1] = deriv_end * dx[n - 2];
s_[n - 1] = deriv_end * dx[n - 2];
break;
case BoundaryCondition::second_deriv:
diag[n - 1] = 2.0 * dx[n - 2];
subdiag[n - 2] = 1.0 * dx[n - 2];
s[n - 1] = (3.0 * dd[n - 2] + 0.5 * deriv_end * dx[n - 2]) * dx[n - 2];
s_[n - 1] = (3.0 * dd[n - 2] + 0.5 * deriv_end * dx[n - 2]) * dx[n - 2];
break;
default:
assert(false && "CubicSpline: BoundaryCondition error!"); // not supposed to happend
default: // BoundaryCondition::not_a_knot
diag[n - 1] = dx[n - 3];
subdiag[n - 2] = x[n - 1] - x[n - 3];
s_[n - 1] = (dd[n - 2] * dx[n - 3] * (dx[n - 2] + 2 * (x[n - 1] - x[n - 3]))
+ dd[n - 3] * dx[n - 2] * dx[n - 2])
/ (x[n - 1] - x[n - 3]);
}

int NRHS = 1;
int LDB = n;
int INFO = 0;
int N = n;

dgtsv_(&N, &NRHS, subdiag, diag, supdiag, s, &LDB, &INFO);
}

std::memcpy(c0_, y, sizeof(double) * (n - 1));
std::memcpy(c1_, s, sizeof(double) * (n - 1));
for (int i = 0; i != n - 1; ++i)
{
c3_[i] = (s[i] + s[i + 1] - 2 * dd[i]) / (dx[i] * dx[i]);
c2_[i] = (dd[i] - s[i]) / dx[i] - c3_[i] * dx[i];
dgtsv_(&N, &NRHS, subdiag, diag, supdiag, s_, &LDB, &INFO);
}

delete[] diag;
delete[] subdiag;
delete[] supdiag;
delete[] dx;
delete[] dd;
delete[] s;
}
}

void CubicSpline::get(const int n, const double* const x, double* const y)
void CubicSpline::interp(const int n, const double* const x, double* const y, double* const dy)
{
assert(x_ && c0_ && c1_ && c2_ && c3_); // make sure the interpolant exists

int p = 0;
double w = 0.0;
assert(x_ && y_ && s_); // make sure the interpolant exists
assert(y || dy); // make sure at least one of y or dy is not null

for (int i = 0; i != n; ++i)
{
assert(x[i] >= x_[0] && x[i] <= x_[n_ - 1]);
}
// check that x is in the range of the interpolant
assert(std::all_of(x, x + n, [this](double x) -> bool { return x >= x_[0] && x <= x_[n_ - 1]; }));

std::function<int(double)> search;
if (is_uniform_)
{
double dx = x_[1] - x_[0];
for (int i = 0; i != n; ++i)
{
p = x[i] == x_[n_ - 1] ? n_ - 2 : x[i] / dx;
w = x[i] - x_[p];
y[i] = ((c3_[p] * w + c2_[p]) * w + c1_[p]) * w + c0_[p];
}
search = [this, dx](double x) -> int { return x == x_[n_ - 1] ? n_ - 2 : x / dx; };
}
else
{
for (int i = 0; i != n; ++i)
{
int q = n_ - 1;
while (q - p > 1)
{
int m = (p + q) / 2;
if (x_[m] > x[i])
{
q = m;
}
else
{
p = m;
}
}
w = x[i] - x_[p];
y[i] = ((c3_[p] * w + c2_[p]) * w + c1_[p]) * w + c0_[p];
}
search = [this](double x) -> int { return (std::upper_bound(x_, x_ + n_, x) - x_) - 1; };
}

void (*eval)(double, double, double, double, double, double*, double*) = y ? (dy ? _eval_y_dy : _eval_y) : _eval_dy;

for (int i = 0; i != n; ++i)
{
int p = search(x[i]);
double w = x[i] - x_[p];

double dx = x_[p + 1] - x_[p];
double dd = (y_[p + 1] - y_[p]) / dx;

double c0 = y_[p];
double c1 = s_[p];
double c3 = (s_[p] + s_[p + 1] - 2.0 * dd) / (dx * dx);
double c2 = (dd - s_[p]) / dx - c3 * dx;

eval(w, c0, c1, c2, c3, y + i, dy + i);
}
}

Expand Down
61 changes: 44 additions & 17 deletions source/module_base/cubic_spline.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,10 @@ namespace ModuleBase
* CubicSpline::BoundaryCondition::periodic );
*
* // calculate the values of interpolant at x_interp[i]
* cubspl.get(n_interp, x_interp, y_interp);
* cubspl.interp(n_interp, x_interp, y_interp);
*
* // calculate the values & derivatives of interpolant at x_interp[i]
* cubspl.interp(n_interp, x_interp, y_interp, dy_interp);
* */
class CubicSpline
{
Expand Down Expand Up @@ -85,32 +88,40 @@ class CubicSpline
//! Calculates the values of interpolant.
/*!
* This function evaluates the interpolant at x_interp[i].
* On finish, values are placed in y_interp.
* On finish, interpolated values are placed in y_interp,
* and the derivatives at x_interp[i] are placed in dy_interp.
*
* If y_interp or dy_interp is nullptr, the corresponding values are
* not calculated. y_interp and dy_interp must not be both nullptr.
* */
void get(const int n, //!< [in] number of points to evaluate the interpolant
const double* const x_interp, //!< [in] places where the interpolant is evaluated;
//!< must be within [x_[0], x_[n-1]]
double* const y_interp //!< [out] interpolated values
void interp(const int n, //!< [in] number of points to evaluate the interpolant
const double* const x_interp, //!< [in] places where the interpolant is evaluated;
//!< must be within [x_[0], x_[n-1]]
double* const y_interp, //!< [out] interpolated values
double* const dy_interp = nullptr //!< [out] derivatives at x_interp
);

/// knots of the interpolant
const double* x() const { return x_; }

/// values at knots
const double* y() const { return y_; }

/// first-order derivatives at knots
const double* s() const { return s_; }

private:
//! number of data points
int n_ = 0;

//! knots (x coordinates of data points)
double* x_ = nullptr;

/*! @name Polynomial coefficients of the interpolant.
*
* The i-th piece polynomial P[i](x) (i=0,1,...) is given by
* P[i](x) = c0[i] + c1[i]*(x-x[i]) + c2[i]*(x-x[i])^2 + c3[i]*(x-x[i])^3
* */
///@{
double* c0_ = nullptr;
double* c1_ = nullptr;
double* c2_ = nullptr;
double* c3_ = nullptr;
///@}
//! values at knots
double* y_ = nullptr;

//! first-order derivatives at knots
double* s_ = nullptr;

//! A flag that tells whether the knots are evenly spaced.
bool is_uniform_ = false;
Expand Down Expand Up @@ -159,6 +170,22 @@ class CubicSpline

//! Wipes off the interpolant (if any) and deallocates memories.
void cleanup();

static void _eval_y(double w, double c0, double c1, double c2, double c3, double* y, double*)
{
*y = ((c3 * w + c2) * w + c1) * w + c0;
}

static void _eval_dy(double w, double, double c1, double c2, double c3, double*, double* dy)
{
*dy = (3.0 * c3 * w + 2.0 * c2) * w + c1;
}

static void _eval_y_dy(double w, double c0, double c1, double c2, double c3, double* y, double* dy)
{
*y = ((c3 * w + c2) * w + c1) * w + c0;
*dy = (3.0 * c3 * w + 2.0 * c2) * w + c1;
}
};

}; // namespace ModuleBase
Expand Down
Loading