Examples¶

// SPDX-License-Identifier: Apache-2.0

// First steps: tensors with shapes known at compile time or at run time, element-wise expressions
// that compute nothing until assigned, and a result made in one pass.

import strata;
import std;

namespace {

using strata::Dyn;
using strata::Matrix;
using strata::Shape;
using strata::Tensor;

template <typename L, typename R>
concept addable = requires(const L& lhs, const R& rhs) { lhs + rhs; };

}   // namespace

auto main() -> int {
    // A static shape, 2 x 3, is part of the type: a small tensor keeps its elements inside itself.
    // Nested lists are written as the matrix is read, row by row.
    const Matrix<float, 2, 3> a{{0, 1, 2}, {10, 11, 12}};

    // A dynamic shape is given at run time; the elements are on the heap.
    const Matrix<float, Dyn, Dyn> b{Shape<Dyn, Dyn>{2, 3}, 0.5F};

    // An expression only records what to compute, over views of `a` and `b`...
    const auto twice_a_plus_b = (a * 2.0F) + b;

    // ... until a tensor is made of it, in one pass over memory. The type follows the expression:
    // float values, and the static shape both operands agree on.
    Matrix c = twice_a_plus_b;
    static_assert(std::same_as<decltype(c), Tensor<float, Shape<2, 3>>>);

    // Tensors, views and expressions print with std::format, as NumPy prints arrays: over lines,
    // columns aligned; the spec after a second colon is the elements', as for a range.
    std::println("a:\n{}", a);
    std::println("b:\n{::.1f}", b);
    std::println("c = 2a + b:\n{::.1f}", c);

    // An expression prints too, read where it is printed: on one line with `l`, with its kind,
    // element type and shape first with `?`.
    std::println("{:?l}", twice_a_plus_b);

    // Assigning to a tensor of the same extents writes in place, no allocation. Every operator of the
    // elements is there: `c -= b` is `c = c - b`.
    c -= b;
    std::println("c -= b:\n{::.1f}", c);

    // Shapes are checked: static extents that differ do not compile, dynamic ones are checked when
    // the expression is made.
    static_assert(!addable<Tensor<float, Shape<2, 3>>, Tensor<float, Shape<3, 2>>>);
    static_assert(addable<Tensor<float, Shape<2, 3>>, Tensor<float, Shape<Dyn, Dyn>>>);
}

// SPDX-License-Identifier: Apache-2.0

// Generators: tensors with values but no memory. A constant, the identity, a run of numbers, any
// function of the index. They mix into expressions like views, are sliced like views, and become
// tensors when assigned; nothing is stored until then.

import strata;
import std;

namespace {

using strata::block;
using strata::Dyn;
using strata::permute;
using strata::seq;
using strata::Shape;
using strata::slice;
using strata::Tensor;

// A generator prints like any tensor, its values computed as they are printed.
auto print(std::string_view title, const auto& matrix) -> void { std::println("{}:\n{::.2f}", title, matrix); }

auto print_vector(std::string_view title, const auto& vector) -> void {
    std::println("{:>18}: {:l:.2f}", title, vector);
}

}   // namespace

auto main() -> int {
    // Constants: the same value everywhere, of any shape.
    print("zeros<float>(2 x 3)", strata::zeros<float>(Shape<2, 3>{}));
    print("full(2 x 3, 7.5)", strata::full(Shape<Dyn, Dyn>{2, 3}, 7.5));

    // Runs of numbers: iota by a step, linspace between two ends, both ends included.
    print_vector("iota<float>(5)", strata::iota<float>(5));
    print_vector("iota(5, 1, 0.25)", strata::iota<double>(5, 1.0, 0.25));
    print_vector("linspace(0, 1, 5)", strata::linspace(0.0, 1.0, 5));

    // The identity, in any rank, and a function of the index for anything else.
    print("eye<float>(4 x 4)", strata::eye<float>(Shape<4, 4>{}));
    const auto distance = strata::generate(Shape<Dyn, Dyn>{4, 5}, [](std::size_t i, std::size_t j) {
        const auto di = static_cast<double>(static_cast<std::ptrdiff_t>(i));
        const auto dj = static_cast<double>(static_cast<std::ptrdiff_t>(j));
        return std::hypot(di, dj);
    });
    print("distance from (0, 0)", distance);

    // Slicing a generator gives a generator: the function is fed other indices, nothing is computed
    // or stored. A row of the identity, and a block of it off the diagonal: a `seq<3>{1}` is three
    // indices from 1, a `seq{1, 4}` the indices 1 to 3 when the size is known only at run time.
    const auto eye4 = strata::eye<float>(Shape<4, 4>{});
    print_vector("eye(4 x 4) row 2", eye4 | slice<0>(2));
    print("eye(4 x 4) 3 x 3 from (0, 1)", eye4 | block(seq<3>{0}, seq<3>{1}));
    print("distance transposed, cut", distance | permute<1, 0> | block(seq{1, 4}, seq{0, 3}));

    // In expressions, like any operand: a matrix weighted column by column, 0 to 1, without a tensor
    // of weights. The weights are a generator too, read by another.
    const auto data = strata::full(Shape<Dyn, Dyn>{3, 4}, 2.0);
    const auto column_weight = strata::generate(
        Shape<Dyn, Dyn>{3, 4},
        [weights = strata::linspace(0.0, 1.0, 4)](std::size_t /*i*/, std::size_t j) { return weights[j]; });
    const Tensor weighted = data * column_weight;
    print("2 weighted by column, 0 .. 1", weighted);

    // Shaped like something at hand, a tensor here, whatever its type.
    const Tensor identity = strata::eye<double>(Shape<3, 3>{});
    const Tensor shifted = identity + strata::full_like(identity, 0.5);
    print("identity + full_like(identity, 0.5)", shifted);
}

// SPDX-License-Identifier: Apache-2.0

// Views: memory someone else owns, seen as a tensor. Slices, blocks and transposes are views too,
// of the same memory: nothing is copied, and writing through them writes the original.

import strata;
import std;

namespace {

using strata::block;
using strata::Dense;
using strata::Dyn;
using strata::permute;
using strata::RowMajor;
using strata::seq;
using strata::Shape;
using strata::slice;
using strata::Strided;
using strata::Tensor;
using strata::TensorView;

// A function over any matrix of doubles takes a strided view: every layout converts to it, a
// tensor, a row-major view, a block, a transpose.
using AnyMatrix = TensorView<const double, Shape<Dyn, Dyn>, Strided<Dyn, Dyn>>;

auto print(std::string_view title, AnyMatrix matrix) -> void {
    std::println("{} {}:\n{}", title, matrix.shape(), matrix);
}

// The sum of the diagonal, of a square matrix or not.
auto trace(AnyMatrix matrix) -> double {
    double total = 0;
    for (std::size_t i = 0; i < std::min(matrix.extent(0), matrix.extent(1)); ++i) {
        total += matrix[i, i];
    }
    return total;
}

}   // namespace

auto main() -> int {
    // Twelve numbers in a std::vector, filled through a flat view of it, then seen as a row-major
    // 3 x 4 matrix, as C and NumPy lay it out.
    std::vector<double> numbers(12);
    strata::assign(TensorView{numbers.data(), Shape<Dyn>{numbers.size()}}, strata::iota<double>(12));
    const TensorView<double, Shape<Dyn, Dyn>, Dense<RowMajor>> matrix{numbers.data(), Shape<Dyn, Dyn>{3, 4}};
    print("matrix", matrix);

    // Views are made with a pipe: `matrix | slice<0>(1)` is row 1, an axis less. The column steps over
    // the rows, so its view is strided.
    const auto row = matrix | slice<0>(1);
    const auto column = matrix | slice<1>(2);
    std::println("row 1: {:l}", row);
    std::println("column 2: {:l}", column);

    // A block is a `seq` for each axis: `seq<2>{1}` is two indices from 1, so the size is known at compile
    // time; `seq{0, 3}` is the indices 0 to 2, a size known at run time.
    const auto middle = matrix | block(seq<2>{1}, seq<2>{1});
    print("2 x 2 from (1, 1)", middle);
    print("rows 0 to 2, two columns from 2", matrix | block(seq{0, 3}, seq<2>{2}));

    // A transpose reorders the axes; the memory stays where it is.
    print("transposed", matrix | permute<1, 0>);

    // Assigning to a view writes the memory it views, here the vector: the top-left block becomes
    // ten times the bottom-right one. The operators of the elements are all there, `+=` to `>>=`.
    const auto top_left = matrix | block(seq<2>{0}, seq<2>{0});
    const auto bottom_right = matrix | block(seq<2>{1}, seq<2>{2});
    top_left = bottom_right * 10.0;
    row += 100.0;   // row 1 was a view of `numbers` before the write, and still is
    std::println("numbers[0] = {}, numbers[4] = {}, numbers[5] = {}", numbers[0], numbers[4], numbers[5]);
    print("after the writes", matrix);

    // Expressions over views of different layouts mix freely.
    const auto corner = matrix | block(seq<3>{0}, seq<3>{0});
    const Tensor<double, Shape<Dyn, Dyn>> symmetric = corner + (corner | permute<1, 0>);
    print("block + its transpose", symmetric);

    // One function of our own for all of them.
    std::println("traces: matrix {}, block {}, transpose {}, tensor {}", trace(matrix), trace(middle),
                 trace(matrix | permute<1, 0>), trace(symmetric));

    // strata's reductions take them all as well, and read each in the order of its own memory: the
    // row-major matrix front to back, its transpose too, the block row by row.
    std::println("sums: matrix {}, block {}, transpose {}, tensor {}", strata::sum(matrix), strata::sum(middle),
                 strata::sum(matrix | permute<1, 0>), strata::sum(symmetric));
    std::println("matrix: min {}, max {}, mean {}", strata::min(matrix), strata::max(matrix), strata::mean(matrix));

    // A std::array works the same; static extents make a static view.
    std::array<double, 6> raw{1, 2, 3, 4, 5, 6};
    const TensorView fixed{raw.data(), Shape<2, 3>{}};
    print("std::array as 2 x 3", fixed);
}

// SPDX-License-Identifier: Apache-2.0

// Where a tensor's elements live, and how to choose: the default policy, a stronger alignment, an
// arena, and an allocator of one's own. Under a policy strata places the memory itself; an
// allocator given is always used.

import strata;
import std;

namespace {

using strata::Aligned;
using strata::ColumnMajor;
using strata::Dyn;
using strata::Shape;
using strata::Tensor;

// Whether `pointer` is on a `boundary`-byte boundary: what an alignment promises, and no more.
auto aligned_to(const void* pointer, std::size_t boundary) -> bool {
    return std::bit_cast<std::uintptr_t>(pointer) % boundary == 0;
}

// An allocator of one's own: std::allocator underneath, counting what it is asked for. It declares
// the alignment it gives, which is all strata needs to know of it.
std::size_t allocations = 0;

template <typename T>
struct CountingAllocator {
    using value_type = T;

    static constexpr std::size_t alignment = __STDCPP_DEFAULT_NEW_ALIGNMENT__;

    CountingAllocator() = default;

    template <typename U>
    CountingAllocator(const CountingAllocator<U>& /*other*/) noexcept {}

    auto allocate(std::size_t count) -> T* {
        ++allocations;
        return std::allocator<T>{}.allocate(count);
    }

    auto deallocate(T* pointer, std::size_t count) noexcept -> void { std::allocator<T>{}.deallocate(pointer, count); }

    friend auto operator==(const CountingAllocator&, const CountingAllocator&) -> bool = default;
};

}   // namespace

auto main() -> int {
    // The default policy, Aligned<64>: a small static tensor is its elements, a cache line of them;
    // a dynamic one is on the heap, 64-byte aligned for SIMD and cache lines.
    const Tensor<float, Shape<4, 4>> small;
    const Tensor<float, Shape<Dyn, Dyn>> large{Shape<Dyn, Dyn>{100, 100}};
    std::println("Tensor<float, Shape<4, 4>>: {} bytes, inside the object", sizeof(small));
    std::println("Tensor<float, Shape<Dyn, Dyn>>: {} bytes of handle, elements 64-byte aligned: {}", sizeof(large),
                 aligned_to(large.data(), 64));

    // A stronger alignment is a policy too, for inline and heap elements alike.
    const Tensor<float, Shape<8, 8>, ColumnMajor, Aligned<128>> aligned;
    std::println("Aligned<128>: elements 128-byte aligned: {}", aligned_to(aligned.data(), 128));

    // Elements about to be overwritten need no zeros first. A tensor made of an expression skips
    // them by itself; written by hand, it is asked for.
    Tensor<float, Shape<Dyn>> buffer{Shape<Dyn>{1000}, strata::uninitialized};
    strata::assign(buffer, strata::iota<float>(1000));
    std::println("uninitialized, then written: buffer[999] = {}", buffer[999]);

    // An arena: every tensor takes its elements from the buffer below, static ones too, and all of
    // it goes at once when the arena does. A std::pmr::polymorphic_allocator<> serves any element
    // type: strata rebinds it.
    std::array<std::byte, 16 * 1024> stack{};
    {
        std::pmr::monotonic_buffer_resource arena{stack.data(), stack.size(), std::pmr::null_memory_resource()};
        const strata::pmr::Tensor<float, Shape<Dyn, Dyn>> x{Shape<Dyn, Dyn>{16, 16}, 1.0F, &arena};
        const strata::pmr::Tensor<float, Shape<4, 4>> y{strata::eye<float>(Shape<4, 4>{}), &arena};
        const Tensor<double, Shape<Dyn, Dyn>, ColumnMajor, std::pmr::polymorphic_allocator<>> z{x * 3.0F, &arena};
        const auto inside = [&](const void* pointer) {
            const auto* byte = static_cast<const std::byte*>(pointer);
            return byte >= stack.data() && byte < stack.data() + stack.size();
        };
        std::println("arena: x {}, y {}, z {}; y[2, 2] = {}, z[15, 15] = {}", inside(x.data()), inside(y.data()),
                     inside(z.data()), y[2, 2], z[15, 15]);
    }

    // An allocator of one's own is used as std::vector would use it, for every tensor of it.
    Tensor<float, Shape<Dyn>, ColumnMajor, CountingAllocator<float>> counted{Shape<Dyn>{10}, 1.0F};
    counted = counted * 2.0F;   // the same extents: in place
    std::println("counting allocator: {} allocation(s), counted[9] = {}", allocations, counted[9]);
    std::println("strata assumes {}-byte alignment of it", strata::allocation_alignment<CountingAllocator<float>>);
}

// SPDX-License-Identifier: Apache-2.0

// Tensors of quantities with units. strata knows nothing about units: it takes any value type with
// the operations an expression uses, works out the type of each result from them, and has no
// operator where they have none. A units library such as mp-units fits in the same way; a small
// one of our own here keeps the example free of dependencies.

import strata;
import std;

namespace {

using strata::block;
using strata::Dyn;
using strata::seq;
using strata::Shape;
using strata::Tensor;

// A quantity of length^L x time^T, in metres and seconds. The dimension is in the type, so the
// compiler checks it: lengths add to lengths only, and multiplying or dividing makes new ones.
template <int L, int T>
struct Quantity {
    double value;

    friend constexpr auto operator+(Quantity lhs, Quantity rhs) -> Quantity { return {lhs.value + rhs.value}; }
    friend constexpr auto operator-(Quantity lhs, Quantity rhs) -> Quantity { return {lhs.value - rhs.value}; }
    friend constexpr auto operator*(Quantity lhs, double factor) -> Quantity { return {lhs.value * factor}; }

    template <int L2, int T2>
    friend constexpr auto operator*(Quantity lhs, Quantity<L2, T2> rhs) -> Quantity<L + L2, T + T2> {
        return {lhs.value * rhs.value};
    }

    template <int L2, int T2>
    friend constexpr auto operator/(Quantity lhs, Quantity<L2, T2> rhs) -> Quantity<L - L2, T - T2> {
        return {lhs.value / rhs.value};
    }

    // Quantities of one dimension compare; of different ones they do not.
    friend constexpr auto operator<=>(Quantity lhs, Quantity rhs) = default;

    // The root of an even power: of an area a length, of a time squared a time. strata's sqrt calls
    // it unqualified, and argument-dependent lookup finds it here.
    friend auto sqrt(Quantity quantity) -> Quantity<L / 2, T / 2>
        requires(L % 2 == 0 && T % 2 == 0)
    {
        return {std::sqrt(quantity.value)};
    }
};

using Length = Quantity<1, 0>;
using Time = Quantity<0, 1>;
using Speed = Quantity<1, -1>;
using Acceleration = Quantity<1, -2>;

template <int L, int T>
auto unit(Quantity<L, T> /*quantity*/) -> std::string {
    const auto power = [](std::string_view symbol, int exponent) -> std::string {
        if (exponent == 0) {
            return "";
        }
        return exponent == 1 ? std::string{symbol} : std::format("{}^{}", symbol, exponent);
    };
    return power("m", L) + (L != 0 && T != 0 ? "·" : "") + power("s", T);
}

// A tensor of quantities prints once its element does: one line, two decimals each.
auto print(std::string_view title, const auto& quantities) -> void {
    std::println("{:>14}: {:l:.2f}", title, quantities);
}

template <typename L, typename R>
concept addable = requires(const L& lhs, const R& rhs) { lhs + rhs; };

template <typename L, typename R>
concept comparable = requires(const L& lhs, const R& rhs) { lhs < rhs; };

template <typename E>
concept sqrtable = requires(const E& operand) { strata::sqrt(operand); };

}   // namespace

// A quantity prints as its value, by the spec given for a double, and its unit: `{:.2f}` is
// `1.23 m`. With this, a tensor of quantities prints like a tensor of numbers.
template <int L, int T>
struct std::formatter<Quantity<L, T>> : std::formatter<double> {   // NOLINT(bugprone-std-namespace-modification)
    // Any context, not std::format_context only: std::formattable asks with one of its own.
    template <typename Ctx>
    auto format(Quantity<L, T> quantity, Ctx& ctx) const -> Ctx::iterator {
        auto out = std::formatter<double>::format(quantity.value, ctx);
        return std::format_to(out, " {}", unit(quantity));
    }
};

auto main() -> int {
    // Positions of a body falling from rest, every half second: x = g t² / 2. A generator of
    // lengths, made into a tensor.
    constexpr std::size_t samples = 7;
    constexpr Time step{0.5};
    const Tensor<Length, Shape<Dyn>> position = strata::generate(Shape<Dyn>{samples}, [](std::size_t i) {
        const double t = step.value * static_cast<double>(static_cast<std::ptrdiff_t>(i));
        return Length{9.81 * t * t / 2.0};
    });

    // Finite differences over blocks of the same tensor, one step apart. Length / Time is Speed,
    // Speed / Time is Acceleration: the types say so, and so would the compiler if they did not.
    const auto later = position | block(seq{1, samples});
    const auto earlier = position | block(seq{0, samples - 1});
    const Tensor velocity = (later - earlier) / step;
    const Tensor acceleration =
        ((velocity | block(seq{1, samples - 1})) - (velocity | block(seq{0, samples - 2}))) / step;
    static_assert(std::same_as<decltype(velocity)::value_type, Speed>);
    static_assert(std::same_as<decltype(acceleration)::value_type, Acceleration>);

    print("position", position);
    print("velocity", velocity);
    print("acceleration", acceleration);

    // A quantity is a scalar like any number: how far each average speed carries in two seconds.
    const Tensor<Length, Shape<Dyn>> distance = velocity * Time{2.0};
    print("in 2 s", distance);

    // Functions and comparisons take quantities too, with the dimensions checked as for operators:
    // t = sqrt(2 x / g) is a time again, and only speeds compare with a speed.
    constexpr Acceleration g{9.81};
    const Tensor<Time, Shape<Dyn>> time = ((position * 2.0) / g) | strata::sqrt;
    print("sqrt(2 x / g)", time);
    std::println("{:>14}: {} of {}", "faster than 10", strata::count(velocity > Speed{10.0}), velocity.extent(0));
    static_assert(!sqrtable<Tensor<Length, Shape<Dyn>>>);   // a length has no root in metres
    static_assert(!comparable<Tensor<Speed, Shape<Dyn>>, Length>);

    // And a constant of quantities is a generator like any other.
    print("start + 1 m", position + strata::full(Shape<Dyn>{samples}, Length{1.0}));

    // Wrong dimensions are no expression at all.
    static_assert(addable<Tensor<Length, Shape<Dyn>>, Tensor<Length, Shape<Dyn>>>);
    static_assert(!addable<Tensor<Length, Shape<Dyn>>, Tensor<Time, Shape<Dyn>>>);
    static_assert(!addable<Tensor<Speed, Shape<Dyn>>, Length>);
}

// SPDX-License-Identifier: Apache-2.0

// An image as a tensor: height x width x channels, row-major, as image files and libraries lay it
// out. Channels are slices, a crop is a block, a grey version is an expression over the channels,
// a threshold is a mask, and converting it back to bytes is an explicit conversion. The pictures
// are drawn in the terminal in 24-bit colour; with NO_COLOR set (https://no-color.org), as
// characters from dark to light.

import strata;
import std;

namespace {

using strata::block;
using strata::Dyn;
using strata::permute;
using strata::RowMajor;
using strata::seq;
using strata::Shape;
using strata::slice;
using strata::Tensor;

using Image = Tensor<std::uint8_t, Shape<Dyn, Dyn, 3>, RowMajor>;
using Grey = Tensor<float, Shape<Dyn, Dyn>, RowMajor>;
using Colour = std::array<std::uint8_t, 3>;

auto colour_wanted() -> bool {
    static const bool wanted = std::getenv("NO_COLOR") == nullptr;   // NOLINT(concurrency-mt-unsafe): read once
    return wanted;
}

auto to_byte(double value) -> std::uint8_t { return static_cast<std::uint8_t>(std::clamp(value, 0.0, 255.0)); }

// Draws `height` x `width` pixels, `pixel(y, x)` giving each colour. In colour, two pixels to a
// character: the upper half block takes the top one as its colour and the bottom one as its
// background. Otherwise one character a pixel, by its luma.
auto draw(std::string_view title, std::size_t height, std::size_t width, const auto& pixel) -> void {
    std::println("{} ({} x {}):", title, height, width);
    if (!colour_wanted()) {
        constexpr std::string_view ramp = " .:-=+*#%@";
        for (std::size_t y = 0; y < height; ++y) {
            for (std::size_t x = 0; x < width; ++x) {
                const Colour c = pixel(y, x);
                const double luma = (0.299 * c[0]) + (0.587 * c[1]) + (0.114 * c[2]);
                std::print("{}", ramp[static_cast<std::size_t>(luma) * (ramp.size() - 1) / 255]);
            }
            std::println("");
        }
        return;
    }
    for (std::size_t y = 0; y < height; y += 2) {
        for (std::size_t x = 0; x < width; ++x) {
            const Colour top = pixel(y, x);
            std::print("\x1b[38;2;{};{};{}m", top[0], top[1], top[2]);
            if (y + 1 < height) {
                const Colour bottom = pixel(y + 1, x);
                std::print("\x1b[48;2;{};{};{}m", bottom[0], bottom[1], bottom[2]);
            } else {
                std::print("\x1b[49m");   // an odd last row: the terminal's own background below
            }
            std::print("▀");
        }
        std::println("\x1b[0m");
    }
}

// An image, its channels, and grey matrices of 0..255, drawn alike.
auto draw_rgb(std::string_view title, const auto& image) -> void {
    draw(title, image.extent(0), image.extent(1),
         [&](std::size_t y, std::size_t x) { return Colour{image[y, x, 0], image[y, x, 1], image[y, x, 2]}; });
}

auto draw_channel(std::string_view title, const auto& channel, std::size_t which) -> void {
    draw(title, channel.extent(0), channel.extent(1), [&](std::size_t y, std::size_t x) {
        Colour colour{};
        colour[which] = static_cast<std::uint8_t>(channel[y, x]);
        return colour;
    });
}

auto draw_grey(std::string_view title, const auto& grey) -> void {
    draw(title, grey.extent(0), grey.extent(1), [&](std::size_t y, std::size_t x) {
        const std::uint8_t level = to_byte(static_cast<double>(grey[y, x]));
        return Colour{level, level, level};
    });
}

// A sunset over the sea, made by one function of the pixel and the channel: a sky from deep blue
// to orange, a sun on the horizon, and the sea below with the sun's light on it.
auto sunset(std::size_t height, std::size_t width) -> Image {
    const auto pixel = [=](std::size_t y, std::size_t x, std::size_t channel) -> std::uint8_t {
        const double u = static_cast<double>(x) / static_cast<double>(width - 1);    // 0 .. 1, left to right
        const double v = static_cast<double>(y) / static_cast<double>(height - 1);   // 0 .. 1, top to bottom
        constexpr double horizon = 0.62;
        const double sun_x = 0.5;
        const double sun_y = horizon - 0.04;
        const double dx = (u - sun_x) * static_cast<double>(width) / static_cast<double>(height);
        const double dy = v - sun_y;
        const double sun = std::exp(-((dx * dx) + (dy * dy)) / 0.006);   // a soft disc
        std::array<double, 3> colour{};
        if (v < horizon) {
            const double t = v / horizon;   // 0 at the top of the sky, 1 at the horizon
            colour = {20 + (235 * t * t), 30 + (110 * t * t), 90 + (20 * t) - (60 * t * t)};
            const double glow = sun + (0.35 * std::exp(-((dx * dx) + (dy * dy)) / 0.05));
            colour = {colour[0] + (255 * glow), colour[1] + (200 * glow), colour[2] + (90 * sun)};
        } else {
            const double depth = (v - horizon) / (1 - horizon);   // 0 at the horizon, 1 at the bottom
            colour = {15 + (40 * (1 - depth)), 40 + (40 * (1 - depth)), 70 + (60 * (1 - depth))};
            const double path = std::exp(-(dx * dx) / (0.004 + (0.05 * depth))) * (1 - (0.6 * depth));
            const double ripple = 0.5 + (0.5 * std::sin((140.0 * v) + (6.0 * u)));
            colour = {colour[0] + (230 * path * ripple), colour[1] + (140 * path * ripple), colour[2] + (40 * path)};
        }
        return to_byte(colour[channel]);
    };
    return Image{strata::generate(Shape<Dyn, Dyn, 3>{height, width}, pixel)};
}

}   // namespace

auto main() -> int {
    constexpr std::size_t height = 40;
    constexpr std::size_t width = 72;
    const Image image = sunset(height, width);
    draw_rgb("a sunset, 40 x 72 x 3 bytes", image);

    // Each channel is a slice along the last axis: a height x width view, stepping 3 bytes at a time.
    const auto red = image | slice<2>(0);
    const auto green = image | slice<2>(1);
    const auto blue = image | slice<2>(2);
    draw_channel("red channel", red, 0);
    draw_channel("blue channel", blue, 2);

    // Luma as a weighted sum of the channels: bytes times floats are floats.
    const Grey grey = (red * 0.299F) + (green * 0.587F) + (blue * 0.114F);
    draw_grey("grey", grey);

    // A crop is a block, of the colour image or of the grey one; a transpose turns it on its side.
    draw_rgb("the sun, cropped", image | block(seq{14, 34}, seq{20, 52}, seq<3>{0}));
    draw_grey("grey crop, transposed", grey | block(seq{14, 34}, seq{20, 52}) | permute<1, 0>);

    // A threshold is a mask, and where draws it: white where the grey is bright, black elsewhere.
    const auto bright = grey > 200.0F;
    draw_grey("grey > 200", strata::where(bright, 255.0F, 0.0F));
    std::println("bright pixels: {} of {}", strata::count(bright), height * width);

    // A vignette: darker towards the corners, from a generator of weights, on each channel alike.
    const auto vignette = strata::generate(Shape<Dyn, Dyn>{height, width}, [](std::size_t y, std::size_t x) {
        const double dy = (static_cast<double>(static_cast<std::ptrdiff_t>(y)) - 19.5) / 22.0;
        const double dx = (static_cast<double>(static_cast<std::ptrdiff_t>(x)) - 35.5) / 38.0;
        return static_cast<float>(std::max(0.0, 1.0 - (0.9 * ((dx * dx) + (dy * dy)))));
    });
    const Tensor<std::uint8_t, Shape<Dyn, Dyn, 3>, RowMajor> vignetted =
        strata::generate(Shape<Dyn, Dyn, 3>{height, width}, [&](std::size_t y, std::size_t x, std::size_t channel) {
            return to_byte(static_cast<double>(image[y, x, channel]) * static_cast<double>(vignette[y, x]));
        });
    draw_rgb("vignetted", vignetted);

    // Back to bytes: another value type, so the conversion is spelt out.
    const Tensor<std::uint8_t, Shape<Dyn, Dyn>, RowMajor> bytes{grey};
    std::println("grey bytes at the sun and at a corner: {} {}", bytes[22, 36], bytes[0, 0]);
}

// SPDX-License-Identifier: Apache-2.0

// Math and masks on a sampled signal. The functions of <cmath> take expressions and give
// expressions; comparisons give masks of bool, which combine with & | ! and choose values with
// where; all, any and count answer questions about them. Signals are drawn as one line of bars.
// Every function takes its operand also from the left of a pipe: `x | abs | sqrt`; a call without its
// first argument, `clamp(-1, 1)`, waits for it there. Comparisons bind tighter than |, so they are
// parenthesised: `(x | abs) > 1.5`.

import strata;
import std;

namespace {

using strata::Dyn;
using strata::Shape;
using strata::Tensor;

using Signal = Tensor<double, Shape<Dyn>>;
using Mask = Tensor<bool, Shape<Dyn>>;

constexpr std::size_t samples = 64;

// A line of bars, the lowest for `low` and the highest for `high`, clamped between.
auto draw(std::string_view title, const auto& signal, double low = -3.0, double high = 3.0) -> void {
    constexpr std::array<std::string_view, 8> bars{"▁", "▂", "▃", "▄", "▅", "▆", "▇", "█"};
    std::print("{:>26} ", title);
    for (std::size_t i = 0; i < signal.extent(0); ++i) {
        const double level = std::clamp((signal[i] - low) / (high - low), 0.0, 1.0);
        std::print("{}", bars[static_cast<std::size_t>(std::lround(level * 7.0))]);
    }
    std::println("");
}

// A mask as a line of marks: `|` where true.
auto draw_mask(std::string_view title, const auto& mask) -> void {
    std::print("{:>26} ", title);
    for (std::size_t i = 0; i < mask.extent(0); ++i) {
        std::print("{}", mask[i] ? "|" : "·");
    }
    std::println("");
}

}   // namespace

auto main() -> int {
    // A slow wave with a fast ripple on it, and a few spikes: a sensor that glitches now and then.
    const auto t = strata::linspace(0.0, 2.0 * std::numbers::pi, samples);
    const auto spikes = strata::generate(Shape<Dyn>{samples},
                                         [](std::size_t i) { return i % 17 == 5 ? 3.0 : (i % 23 == 11 ? -3.0 : 0.0); });
    const Signal signal = strata::sin(t) + (strata::sin(t * 7.0) * 0.3) + spikes;
    draw("signal", signal);

    // Element-wise math is lazy like any expression: nothing is computed until assigned or read.
    draw("abs(signal)", signal | strata::abs);
    draw("clamp(signal, -1, 1)", signal | strata::clamp(-1.0, 1.0));
    draw("maximum(signal, 0)", signal | strata::maximum(0.0));

    // A comparison is a mask; masks combine element by element. && and || are left as they are:
    // overloaded, they would no longer short-circuit.
    const auto glitch = (signal | strata::abs) > 1.5;
    const auto rising = strata::cos(t) > 0.0;
    draw_mask("abs(signal) > 1.5", glitch);
    draw_mask("cos(t) > 0", rising);
    draw_mask("glitch & cos(t) > 0", glitch & rising);
    draw_mask("!glitch", !glitch);

    // Questions about a mask, all at once over every element.
    std::println("glitches: {} of {}; any: {}; all clean: {}", strata::count(glitch), samples, strata::any(glitch),
                 strata::all(!glitch));

    // where chooses, element by element: the glitches replaced by the wave they hide.
    const Signal cleaned = strata::where(glitch, strata::sin(t), signal);
    draw("where(glitch, sin(t), x)", cleaned);
    std::println("glitches left: {}", ((cleaned | strata::abs) > 1.5) | strata::count);

    // A mask kept as a tensor of bool, and a function of our own over several operands.
    const Mask positive = cleaned > 0.0;
    const auto smoothstep = strata::map(
        [](double x, double edge) {
            const double s = std::clamp(x / edge, 0.0, 1.0);
            return s * s * (3.0 - (2.0 * s));
        },
        cleaned, 0.8);
    draw_mask("cleaned > 0", positive);
    draw("smoothstep(cleaned, 0.8)", smoothstep, 0.0, 1.0);

    // A logarithm of a signal that touches zero is no longer finite; isfinite says where. A pipeline
    // needs no data to be written down: `floored` is a value, to apply to any signal.
    const auto decibels = (cleaned | strata::abs | strata::log10) * 20.0;
    constexpr auto floored = strata::abs | strata::maximum(1e-6) | strata::log10 | strata::isfinite;
    std::println("finite decibels: {} of {}; after a floor: {}", strata::isfinite(decibels) | strata::count, samples,
                 cleaned | floored | strata::all);

    // == stays one bool, as std::equality_comparable expects; equal() is the element-wise one.
    const Signal rounded = strata::round(cleaned);
    std::println("rounded to zero: {}; to one: {}", strata::count(strata::equal(rounded, 0.0)),
                 strata::count(strata::equal(rounded, 1.0)));
}