Skip to content

File fdarray.hpp

File List > include > multigrid > fdarray.hpp

Go to the documentation of this file

#pragma once

#include "multigrid/boundary_conditions.hpp"
#include "multigrid/fdbase.hpp"
#include "multigrid/field.hpp"
#include "multigrid/types.hpp"
#include <cmath>
#include <functional>

namespace mgrid {

// Forward declaration for gradient() (defined in Task 13)
struct FDVecArray;

class FDArray : public FDBase {
  public:
    FDArray() = default;

    FDArray(double aspect, int nx, int nz);

    FDArray(const FDArray &) = default;
    FDArray(FDArray &&) = default;
    FDArray &operator=(FDArray &&) = default;
    ~FDArray() override = default;

    BoundaryConditions boundaryConditions;

    void resize(double aspect, int nx, int nz, AxisKind kx = AxisKind::bounded,
                AxisKind kz = AxisKind::bounded);

    [[nodiscard]] double &operator()(int i, int j) { return field_(i, j); }
    [[nodiscard]] double operator()(int i, int j) const { return field_(i, j); }
    [[nodiscard]] double &operator[](int i, int j) { return field_(i, j); }
    [[nodiscard]] double operator[](int i, int j) const { return field_(i, j); }

    [[nodiscard]] int rows() const noexcept { return field_.rows(); }
    [[nodiscard]] int cols() const noexcept { return field_.cols(); }
    [[nodiscard]] Field2D &field() noexcept { return field_; }
    [[nodiscard]] const Field2D &field() const noexcept { return field_; }
    [[nodiscard]] double *data() noexcept { return field_.data(); }
    [[nodiscard]] const double *data() const noexcept { return field_.data(); }

    FDArray &operator=(double s) {
        field_.fill(s);
        return *this;
    }
    FDArray &operator=(const Field2D &v) {
        field_ = v;
        return *this;
    }
    FDArray &operator=(const FDArray &o);
    FDArray &operator=(const std::function<double(double, double)> &f);

    FDArray &operator+=(double s) {
        field_ += s;
        return *this;
    }
    FDArray &operator-=(double s) {
        field_ -= s;
        return *this;
    }
    FDArray &operator*=(double s) {
        field_ *= s;
        return *this;
    }
    FDArray &operator/=(double s) {
        field_ /= s;
        return *this;
    }
    FDArray &operator+=(const Field2D &v) {
        field_ += v;
        return *this;
    }
    FDArray &operator-=(const Field2D &v) {
        field_ -= v;
        return *this;
    }
    FDArray &operator+=(const FDArray &o) {
        field_ += o.field_;
        return *this;
    }
    FDArray &operator-=(const FDArray &o) {
        field_ -= o.field_;
        return *this;
    }
    FDArray &operator*=(const FDArray &o) {
        field_ *= o.field_;
        return *this;
    }
    FDArray &operator/=(const FDArray &o) {
        field_ /= o.field_;
        return *this;
    }

    [[nodiscard]] double norm() const { return frobenius_norm(field_); }

    // --- finite differences: added in Tasks 9-11 ---

    // clang-format off
    // The one-sided / 9-region stencils below are transcribed term-by-term from
    // the 2010-2011 sources (and the dxz (nx-1, 0) sign-typo correction is
    // annotated inline); their one-coefficient-per-line layout is checked by eye
    // against the Taylor expansions, so `just format` must leave it alone.

    [[nodiscard]] double dx(int i, int j) const {
        const Axis& ax = x_axis();
        if (ax.edge_lo(i)) {
            return (-3 * field_(i, j) + 4 * field_(i + 1, j) - field_(i + 2, j)) * xfactor_;
        }
        if (ax.edge_hi(i)) {
            return (field_(i - 2, j) - 4 * field_(i - 1, j) + 3 * field_(i, j)) * xfactor_;
        }
        return (field_(ax.shift(i, 1), j) - field_(ax.shift(i, -1), j)) * xfactor_;
    }

    [[nodiscard]] double dxu(int i, [[maybe_unused]] int j) const {
        const Axis& ax = x_axis();
        if (ax.edge_lo(i)) { return -3 * xfactor_; }
        if (ax.edge_hi(i)) { return 3 * xfactor_; }
        return 0;
    }

    [[nodiscard]] double dz(int i, int j) const {
        const Axis& az = z_axis();
        if (az.edge_lo(j)) {
            return (-3 * field_(i, j) + 4 * field_(i, j + 1) - field_(i, j + 2)) * zfactor_;
        }
        if (az.edge_hi(j)) {
            return (field_(i, j - 2) - 4 * field_(i, j - 1) + 3 * field_(i, j)) * zfactor_;
        }
        return (field_(i, az.shift(j, 1)) - field_(i, az.shift(j, -1))) * zfactor_;
    }

    [[nodiscard]] double dzu([[maybe_unused]] int i, int j) const {
        const Axis& az = z_axis();
        if (az.edge_lo(j)) { return -3 * zfactor_; }
        if (az.edge_hi(j)) { return 3 * zfactor_; }
        return 0;
    }

    [[nodiscard]] double dxx(int i, int j) const {
        const Axis& ax = x_axis();
        if (ax.edge_lo(i)) {
            return (-field_(i + 3, j) + 4 * field_(i + 2, j)
                    - 5 * field_(i + 1, j) + 2 * field_(i, j)) * xxfactor_;
        }
        if (ax.edge_hi(i)) {
            return (-field_(i - 3, j) + 4 * field_(i - 2, j)
                    - 5 * field_(i - 1, j) + 2 * field_(i, j)) * xxfactor_;
        }
        return (field_(ax.shift(i, -1), j) - 2 * field_(i, j)
                + field_(ax.shift(i, 1), j)) * xxfactor_;
    }

    [[nodiscard]] double dxxu(int i, [[maybe_unused]] int j) const {
        const Axis& ax = x_axis();
        // low/high one-sided 2nd-difference of a constant are equal by construction.
        // NOLINTNEXTLINE(bugprone-branch-clone)
        if (ax.edge_lo(i)) { return 2 * xxfactor_; }
        if (ax.edge_hi(i)) { return 2 * xxfactor_; }
        return -2 * xxfactor_;
    }

    [[nodiscard]] double dzz(int i, int j) const {
        const Axis& az = z_axis();
        if (az.edge_lo(j)) {
            return (-field_(i, j + 3) + 4 * field_(i, j + 2)
                    - 5 * field_(i, j + 1) + 2 * field_(i, j)) * zzfactor_;
        }
        if (az.edge_hi(j)) {
            return (-field_(i, j - 3) + 4 * field_(i, j - 2)
                    - 5 * field_(i, j - 1) + 2 * field_(i, j)) * zzfactor_;
        }
        return (field_(i, az.shift(j, -1)) - 2 * field_(i, j)
                + field_(i, az.shift(j, 1))) * zzfactor_;
    }

    [[nodiscard]] double dzzu([[maybe_unused]] int i, int j) const {
        const Axis& az = z_axis();
        // low/high one-sided 2nd-difference of a constant are equal by construction.
        // NOLINTNEXTLINE(bugprone-branch-clone)
        if (az.edge_lo(j)) { return 2 * zzfactor_; }
        if (az.edge_hi(j)) { return 2 * zzfactor_; }
        return -2 * zzfactor_;
    }

    [[nodiscard]] double dxz(int i, int j) const {
        const Axis& ax = x_axis();
        const Axis& az = z_axis();
        if (az.edge_lo(j)) {
            if (ax.edge_lo(i)) {
                // forward difference in i, forward difference in j
                return (16*field_(i+1,j+1) - 12*(field_(i+1,j) + field_(i,j+1))
                    + 9*field_(i,j) - 4*(field_(i+1,j+2) + field_(i+2,j+1))
                    + 3*(field_(i+2,j) + field_(i,j+2)) + field_(i+2,j+2))*xzfactor_;
            } else if (ax.edge_hi(i)) {
                // backward difference in i, forward difference in j
                return (-16*field_(i-1,j+1) + 12*(field_(i-1,j) + field_(i,j+1))
                    - 9*field_(i,j) + 4*(field_(i-1,j+2) + field_(i-2,j+1))
                    - 3*(field_(i-2,j) + field_(i,j+2)) // legacy fdarray.hpp:269 had '- 3*(A - B)'; corrected to '+' so this corner's 9 stencil coefficients sum to 0, matching the other three corners
                    - field_(i-2,j+2))*xzfactor_;
            } else {
                // centered difference in i, forward difference in j; the centred
                // i-index wraps on a periodic x-axis (identity on a bounded one)
                return (4*(field_(ax.shift(i,1),j+1)-field_(ax.shift(i,-1),j+1))
                    + 3*(field_(ax.shift(i,-1),j)-field_(ax.shift(i,1),j)) + field_(ax.shift(i,-1),j+2)
                    - field_(ax.shift(i,1),j+2))*xzfactor_;
            }
        } else if (az.edge_hi(j)) {
            if (ax.edge_lo(i)) {
                // forward difference in i, backward difference in j
                return (-16*field_(i+1,j-1) + 12*(field_(i+1,j) + field_(i,j-1))
                    - 9*field_(i,j) + 4*(field_(i+2,j-1) + field_(i+1,j-2))
                    - 3*(field_(i+2,j) + field_(i,j-2)) - field_(i+2,j-2))*xzfactor_;
            } else if (ax.edge_hi(i)) {
                // backward difference in i, backward difference in j
                return (16*field_(i-1,j-1) - 12*(field_(i-1,j) + field_(i,j-1))
                    + 9*field_(i,j) - 4*(field_(i-1,j-2) + field_(i-2,j-1))
                    + 3*(field_(i-2,j) + field_(i,j-2)) + field_(i-2,j-2))*xzfactor_;
            } else {
                // centered difference in i, backward difference in j; the centred
                // i-index wraps on a periodic x-axis (identity on a bounded one)
                return (4*(field_(ax.shift(i,-1),j-1) - field_(ax.shift(i,1),j-1))
                    + 3*(field_(ax.shift(i,1),j) - field_(ax.shift(i,-1),j))
                    - field_(ax.shift(i,-1),j-2) + field_(ax.shift(i,1),j-2))*xzfactor_;
            }
        } else {
            if (ax.edge_lo(i)) {
                // forward difference in i, centered difference in j; the centred
                // j-index wraps on a periodic z-axis (identity on a bounded one)
                return (3*(field_(i,az.shift(j,-1)) - field_(i,az.shift(j,1)))
                    - 4*(field_(i+1,az.shift(j,-1)) - field_(i+1,az.shift(j,1)))
                    + field_(i+2,az.shift(j,-1)) - field_(i+2,az.shift(j,1)))*xzfactor_;
            } else if (ax.edge_hi(i)) {
                // backward difference in i, centered difference in j; the centred
                // j-index wraps on a periodic z-axis (identity on a bounded one)
                return (3*(field_(i,az.shift(j,1)) - field_(i,az.shift(j,-1)))
                    - 4*(field_(i-1,az.shift(j,1)) - field_(i-1,az.shift(j,-1)))
                    + field_(i-2,az.shift(j,1)) - field_(i-2,az.shift(j,-1)))*xzfactor_;
            } else {
                // centered difference in i, centered difference in j; neighbour
                // indices wrap on a periodic axis (identity on a bounded one)
                return (field_(ax.shift(i,-1),az.shift(j,-1))
                    - field_(ax.shift(i,-1),az.shift(j,1))
                    - field_(ax.shift(i,1),az.shift(j,-1))
                    + field_(ax.shift(i,1),az.shift(j,1)))*xzfactor_;
            }
        }
    }

    [[nodiscard]] double dxzu(int i, int j) const {
        const Axis& ax = x_axis();
        const Axis& az = z_axis();
        if (az.edge_lo(j)) {
            if (ax.edge_lo(i)) {
                // forward difference in i, forward difference in j
                return 9*xzfactor_;
            } else if (ax.edge_hi(i)) {
                // backward difference in i, forward difference in j
                return -9*xzfactor_;
            } else {
                // centered difference in i, forward difference in j
                return 0;
            }
        } else if (az.edge_hi(j)) {
            if (ax.edge_lo(i)) {
                // forward difference in i, backward difference in j
                return -9*xzfactor_;
            } else if (ax.edge_hi(i)) {
                // backward difference in i, backward difference in j
                return 9*xzfactor_;
            } else {
                // centered difference in i, backward difference in j
                return 0;
            }
        } else {
            // dxzu of a constant field is zero on this whole region (the interior
            // case a periodic axis always hits); the i-branches are kept to
            // mirror dxz's 9-region stencil.
            // NOLINTNEXTLINE(bugprone-branch-clone)
            if (ax.edge_lo(i)) {
                // forward difference in i, centered difference in j
                return 0;
            } else if (ax.edge_hi(i)) {
                // backward difference in i, centered difference in j
                return 0;
            } else {
                // centered difference in i, centered difference in j
                return 0;
            }
        }
    }

    // clang-format on

    // --- whole-array derivatives / gradient / flux: added in Task 11 ---

    void derivatives(Deriv &d, int i, int j) const {
        d.dx = dx(i, j);
        d.dxu = dxu(i, j);
        d.dz = dz(i, j);
        d.dzu = dzu(i, j);
        d.dxx = dxx(i, j);
        d.dxxu = dxxu(i, j);
        d.dzz = dzz(i, j);
        d.dzzu = dzzu(i, j);
        d.dxz = dxz(i, j);
        d.dxzu = dxzu(i, j);
    }

    void dx(Field2D &out) const {
        for_each_index(out, [this, &out](int i, int j) { out(i, j) = dx(i, j); });
    }

    void dz(Field2D &out) const {
        for_each_index(out, [this, &out](int i, int j) { out(i, j) = dz(i, j); });
    }

    void dxx(Field2D &out) const {
        for_each_index(out, [this, &out](int i, int j) { out(i, j) = dxx(i, j); });
    }

    void dzz(Field2D &out) const {
        for_each_index(out, [this, &out](int i, int j) { out(i, j) = dzz(i, j); });
    }

    void dxz(Field2D &out) const {
        for_each_index(out, [this, &out](int i, int j) { out(i, j) = dxz(i, j); });
    }

    void gradient_magnitude(Field2D &out) const {
        for_each_index(out, [this, &out](int i, int j) {
            double dx_val = dx(i, j);
            double dz_val = dz(i, j);
            out(i, j) = std::sqrt(dx_val * dx_val + dz_val * dz_val);
        });
    }

    void gradient(FDVecArray &result) const;

    [[nodiscard]] double calculate_flux() const;

    // --- boundary update: added in Task 12 ---

    void update_boundaries();

    // --- write(): added in Task 22 ---

    virtual void write(std::string filestring);

  private:
    Field2D field_;
};

} // namespace mgrid