yum/3ner

A toon shader for Unity's BIRP.

git clone https://git.yummers.dev/yum/3ner

yumImpostors: begin optimization work50e4514

master
10.7 KiB348 linesraw
1#!/usr/bin/env python3
2"""
3Fourier approximation utility.
4
5Given an analytic expression f(x) and a finite interval [a, b], this script
6samples the function uniformly, computes Fourier series coefficients via the
7FFT, and reports the L2 error of partial sums. Sine and cosine components are
8treated as individual terms and applied in descending order of amplitude to
9highlight the strongest contributions first.
10"""
11
12from __future__ import annotations
13
14import argparse
15import numpy as np
16import sys
17from typing import Callable, Dict, NamedTuple
18
19
20class FourierTerm(NamedTuple):
21    index: int
22    kind: str  # "sin" or "cos"
23    coefficient: float
24    amplitude: float
25    phase: float
26    angular_frequency: float
27    frequency: float
28
29def frac(x):
30    """Return the fractional part of x in [0, 1)."""
31    arr = np.asarray(x, dtype=float)
32    frac_part = arr - np.floor(arr)
33    if np.isscalar(x):
34        return float(frac_part)
35    return frac_part
36
37# Functions that are allowed inside the user supplied expression. The subset is
38# intentionally small to keep evaluation safe while still being practical.
39_ALLOWED_FUNCS: Dict[str, object] = {
40    "np": np,
41    "sin": np.sin,
42    "cos": np.cos,
43    "tan": np.tan,
44    "exp": np.exp,
45    "log": np.log,
46    "log10": np.log10,
47    "log2": np.log2,
48    "sqrt": np.sqrt,
49    "sinh": np.sinh,
50    "cosh": np.cosh,
51    "tanh": np.tanh,
52    "arcsin": np.arcsin,
53    "arccos": np.arccos,
54    "arctan": np.arctan,
55    "abs": np.abs,
56    "pi": np.pi,
57    "e": np.e,
58    "frac": frac,
59}
60
61
62def build_function(expression: str) -> Callable[[np.ndarray], np.ndarray]:
63    """Create a vectorised callable from the provided expression string."""
64    try:
65        code = compile(expression, "<expression>", "eval")
66    except SyntaxError as exc:
67        raise ValueError(f"Invalid function expression: {exc}") from exc
68
69    def func(x: np.ndarray) -> np.ndarray:
70        local_dict = dict(_ALLOWED_FUNCS)
71        local_dict["x"] = x
72        try:
73            value = eval(code, {"__builtins__": {}}, local_dict)
74        except Exception as exc:  # pragma: no cover - user provided expression
75            raise ValueError(f"Error while evaluating expression: {exc}") from exc
76        return np.asarray(value, dtype=float)
77
78    return func
79
80
81def l2_norm(values: np.ndarray, length: float) -> float:
82    """Compute the L2 norm using a midpoint rule over the sampling grid."""
83    squared = np.square(values)
84    step = length / values.size
85    integral = np.sum(squared) * step
86    return float(np.sqrt(integral / length))
87
88
89def fft_terms(
90    func: Callable[[np.ndarray], np.ndarray],
91    interval: tuple[float, float],
92    term_count: int,
93    samples: int,
94) -> tuple[np.ndarray, np.ndarray, float, list[FourierTerm], float, int]:
95    """Sample the function and return Fourier terms up to term_count."""
96    start, end = interval
97    if end <= start:
98        raise ValueError("Interval end must be greater than start.")
99    if term_count < 1:
100        raise ValueError("term_count must be at least 1.")
101    if samples < 2:
102        raise ValueError("samples must be at least 2.")
103
104    length = end - start
105    xs = start + (np.arange(samples) * length / samples)
106
107    fx = func(xs)
108    if fx.shape == ():
109        fx = np.full_like(xs, float(fx))
110    if fx.shape != xs.shape:
111        raise ValueError(
112            "Function evaluation did not return values of the expected shape."
113        )
114
115    spectrum = np.fft.rfft(fx) / samples
116
117    constant_term = float(spectrum[0].real)
118
119    max_terms = min(term_count, max(len(spectrum) - 1, 0))
120    terms: list[FourierTerm] = []
121
122    if max_terms == 0:
123        return xs, fx, constant_term, terms, length, max_terms
124
125    tol = 1e-12
126
127    for n in range(1, max_terms + 1):
128        coeff = spectrum[n]
129        an = float(2.0 * coeff.real)
130        bn = float(-2.0 * coeff.imag)
131        angular_frequency = float(2.0 * np.pi * n / length)
132        frequency = float(n / length)
133
134        if abs(an) > tol:
135            amplitude = abs(an)
136            phase = 0.0 if an >= 0 else float(np.pi)
137            terms.append(
138                FourierTerm(
139                    index=n,
140                    kind="cos",
141                    coefficient=an,
142                    amplitude=amplitude,
143                    phase=phase,
144                    angular_frequency=angular_frequency,
145                    frequency=frequency,
146                )
147            )
148
149        if abs(bn) > tol:
150            amplitude = abs(bn)
151            phase = 0.0 if bn >= 0 else float(np.pi)
152            terms.append(
153                FourierTerm(
154                    index=n,
155                    kind="sin",
156                    coefficient=bn,
157                    amplitude=amplitude,
158                    phase=phase,
159                    angular_frequency=angular_frequency,
160                    frequency=frequency,
161                )
162            )
163
164    return xs, fx, constant_term, terms, length, max_terms
165
166
167def format_partial_expression(
168    constant_term: float,
169    terms: list[FourierTerm],
170    interval_start: float,
171    length: float,
172) -> str:
173    """Build a copy-pastable expression for the partial trigonometric sum."""
174
175    tol = 1e-12
176    expr_parts: list[str] = []
177
178    if abs(constant_term) > tol:
179        expr_parts.append(f"{constant_term:.6e}")
180
181    def append_component(components: list[str], coeff: float, func: str, argument: str) -> None:
182        if abs(coeff) <= tol:
183            return
184        base = f"{abs(coeff):.6e} * {func}({argument})"
185        if components:
186            sign = "+" if coeff >= 0 else "-"
187            components.append(f"{sign} {base}")
188        else:
189            components.append(base if coeff >= 0 else f"-{base}")
190
191    for term in terms:
192        omega = 2.0 * np.pi * term.index / length
193        if abs(interval_start) <= tol:
194            argument = f"{omega:.6e}*x"
195        elif interval_start < 0:
196            argument = f"{omega:.6e}*(x + {abs(interval_start):.6e})"
197        else:
198            argument = f"{omega:.6e}*(x - {interval_start:.6e})"
199        func_name = "sin" if term.kind == "sin" else "cos"
200        append_component(expr_parts, term.coefficient, func_name, argument)
201
202    if not expr_parts:
203        return "0"
204
205    return " ".join(expr_parts)
206
207
208def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
209    parser = argparse.ArgumentParser(
210        description=(
211            "Approximate a real-valued function on [start, end] using a Fourier "
212            "series derived from FFT samples and report the L2 error of the "
213            "first N partial sums (sorted by amplitude)."
214        )
215    )
216    parser.add_argument(
217        "expression",
218        help=(
219            "Function expression in terms of x. Use numpy-style syntax, e.g. "
220            "'sin(x) + 0.5*cos(3*x)'."
221        ),
222    )
223    parser.add_argument(
224        "start",
225        type=float,
226        help="Beginning of the interval for the approximation.",
227    )
228    parser.add_argument(
229        "end",
230        type=float,
231        help="End of the interval for the approximation.",
232    )
233    parser.add_argument(
234        "--terms",
235        type=int,
236        default=10,
237        help="Number of Fourier harmonics to include (default: 10).",
238    )
239    parser.add_argument(
240        "--samples",
241        type=int,
242        default=2048,
243        help="Number of uniform samples for the FFT (default: 2048).",
244    )
245    parser.add_argument(
246        "--relative",
247        action="store_true",
248        help="Report the relative L2 error in addition to the absolute error.",
249    )
250    if argv is None:
251        argv = sys.argv[1:]
252    if not argv:
253        parser.print_help(sys.stderr)
254        parser.exit(1)
255    return parser.parse_args(argv)
256
257
258def main(argv: list[str] | None = None) -> int:
259    args = parse_args(argv)
260    try:
261        func = build_function(args.expression)
262        xs, fx, constant_term, terms, length, available_harmonics = fft_terms(
263            func=func,
264            interval=(args.start, args.end),
265            term_count=args.terms,
266            samples=args.samples,
267        )
268    except ValueError as exc:
269        print(f"Error: {exc}", file=sys.stderr)
270        return 1
271
272    base_norm = l2_norm(fx, length)
273
274    if args.terms > available_harmonics:
275        print(
276            f"Warning: Requested {args.terms} harmonics but only {available_harmonics} available with {args.samples} samples.",
277            file=sys.stderr,
278        )
279
280    theta = 2.0 * np.pi * (xs - args.start) / length
281    sorted_terms = sorted(terms, key=lambda term: abs(term.amplitude), reverse=True)
282    available_components = len(sorted_terms)
283
284    if args.terms > available_components:
285        print(
286            f"Warning: Requested {args.terms} components but only {available_components} available from the sampled harmonics.",
287            file=sys.stderr,
288        )
289
290    sorted_terms = sorted_terms[: args.terms]
291
292    partial = np.full_like(fx, constant_term)
293    cumulative_terms: list[FourierTerm] = []
294    trig_cache: dict[int, tuple[np.ndarray, np.ndarray]] = {}
295
296    header = "Terms".rjust(5) + "  " + "L2 error".rjust(14)
297    if args.relative:
298        header += "  " + "Rel. L2 error".rjust(14)
299    print(header)
300    print("-" * len(header))
301
302    print(f"Constant term: {constant_term:.6e}")
303    print(
304        "Terms are sorted by descending amplitude. Each line is a sine or cosine component with params (amplitude, phase, frequency); phase in radians, frequency in cycles per unit."
305    )
306
307    for idx, term in enumerate(sorted_terms, start=1):
308        cos_sin = trig_cache.get(term.index)
309        if cos_sin is None:
310            angles = term.index * theta
311            cos_sin = (np.cos(angles), np.sin(angles))
312            trig_cache[term.index] = cos_sin
313        cos_n, sin_n = cos_sin
314        if term.kind == "cos":
315            partial += term.coefficient * cos_n
316        else:
317            partial += term.coefficient * sin_n
318        error = l2_norm(fx - partial, length)
319        cumulative_terms.append(term)
320        line = f"{idx:5d}  {error:14.6e}"
321        if args.relative:
322            if base_norm > 0.0:
323                rel_error = error / base_norm
324            else:
325                rel_error = float("nan") if error > 0 else 0.0
326            line += f"  {rel_error:14.6e}"
327        term_info = (
328            f"n={term.index} {term.kind} (coeff {term.coefficient:.6e}, amp {term.amplitude:.6e}, phase {term.phase:.6e}, freq {term.frequency:.6e})"
329        )
330        line += f" {term_info}"
331        print(line)
332        expression = format_partial_expression(
333            constant_term,
334            cumulative_terms,
335            args.start,
336            length,
337        )
338        print(f"        expr: {expression}")
339
340    print(
341        f"Interval length: {length:.6g}, samples: {args.samples}, base L2 norm: {base_norm:.6e}"
342    )
343    print("Note: errors use a midpoint-rule approximation on the sampling grid.")
344    return 0
345
346
347if __name__ == "__main__":
348    sys.exit(main())