Spaces:
Running
Running
| """Objective functions and optimisers, in one or two variables. | |
| Positions, gradients and Hessians are NumPy arrays of length 1 or 2, so each | |
| optimiser is written once for both. | |
| """ | |
| import numpy as np | |
| from sympy import Abs, Expr, acos, asin, atan, cos, exp, lambdify, log, pi, sin, sqrt, symbols, tan | |
| from sympy.parsing.sympy_parser import ( | |
| convert_xor, | |
| implicit_multiplication_application, | |
| parse_expr, | |
| standard_transformations, | |
| ) | |
| # Real, so that e.g. d/dx |x| is sign(x) | |
| VARIABLES = symbols("x y", real=True) | |
| ALLOWED_NAMES = { | |
| "sin": sin, "cos": cos, "tan": tan, "asin": asin, "acos": acos, "atan": atan, | |
| "log": log, "ln": log, "sqrt": sqrt, "abs": Abs, "exp": exp, | |
| "e": exp(1), "pi": pi, "π": pi, | |
| } | |
| def parse_function(expression: str, n_variables: int) -> Expr: | |
| """Parse e.g. "x^2 + 3y^2"; raises ValueError with a message for students.""" | |
| if not expression.strip(): | |
| raise ValueError("Enter a function.") | |
| variables = VARIABLES[:n_variables] | |
| names = {str(v): v for v in VARIABLES} | ALLOWED_NAMES | |
| try: | |
| expr = parse_expr( | |
| expression, | |
| local_dict=names, | |
| transformations=standard_transformations + (implicit_multiplication_application, convert_xor), | |
| ) | |
| except Exception: | |
| raise ValueError("This isn't a function sympy can read. Check the brackets and operators.") | |
| if not isinstance(expr, Expr): | |
| raise ValueError("This isn't a function sympy can read. Check the brackets and operators.") | |
| unknown = expr.free_symbols - set(variables) | |
| if unknown: | |
| allowed = " and ".join(str(v) for v in variables) | |
| raise ValueError(f"Unknown name {', '.join(sorted(map(str, unknown)))}: use only {allowed}.") | |
| return expr | |
| class Objective: | |
| """f with its gradient and Hessian, evaluated at a point or on a grid.""" | |
| def __init__(self, expr: Expr, n_variables: int): | |
| variables = VARIABLES[:n_variables] | |
| self.n = n_variables | |
| self._f = lambdify(variables, expr, modules="numpy") | |
| self._gradient = [lambdify(variables, expr.diff(v), modules="numpy") for v in variables] | |
| self._hessian = [[lambdify(variables, expr.diff(a, b), modules="numpy") for b in variables] for a in variables] | |
| def value(self, p: np.ndarray) -> float: | |
| return float(self._f(*p)) | |
| def gradient(self, p: np.ndarray) -> np.ndarray: | |
| return np.array([float(g(*p)) for g in self._gradient]) | |
| def hessian(self, p: np.ndarray) -> np.ndarray: | |
| return np.array([[float(h(*p)) for h in row] for row in self._hessian]) | |
| def values_on(self, *grid: np.ndarray) -> np.ndarray: | |
| # Constant expressions give a scalar; spread it over the grid | |
| return np.broadcast_to(np.asarray(self._f(*grid), dtype=float), grid[0].shape) | |
| def gradient_on(self, *grid: np.ndarray) -> list[np.ndarray]: | |
| return [np.broadcast_to(np.asarray(g(*grid), dtype=float), grid[0].shape) for g in self._gradient] | |
| class SingularHessian(Exception): | |
| pass | |
| # Each optimiser keeps its own state between steps; step(p) returns the next position. | |
| class GradientDescent: | |
| """Heavy-ball momentum: v ← μv − η∇f(x), x ← x + v.""" | |
| def __init__(self, objective: Objective, learning_rate: float, momentum: float): | |
| self.f, self.lr, self.momentum = objective, learning_rate, momentum | |
| self.v = 0.0 | |
| def step(self, p: np.ndarray) -> np.ndarray: | |
| self.v = self.momentum * self.v - self.lr * self.f.gradient(p) | |
| return p + self.v | |
| class Nesterov: | |
| """Gradient taken at the look-ahead point x + μv.""" | |
| def __init__(self, objective: Objective, learning_rate: float, momentum: float): | |
| self.f, self.lr, self.momentum = objective, learning_rate, momentum | |
| self.v = 0.0 | |
| def step(self, p: np.ndarray) -> np.ndarray: | |
| self.v = self.momentum * self.v - self.lr * self.f.gradient(p + self.momentum * self.v) | |
| return p + self.v | |
| class Newton: | |
| """x ← x − H⁻¹∇f(x).""" | |
| def __init__(self, objective: Objective): | |
| self.f = objective | |
| def step(self, p: np.ndarray) -> np.ndarray: | |
| h, g = self.f.hessian(p), self.f.gradient(p) | |
| try: | |
| delta = np.linalg.solve(h, g) | |
| except np.linalg.LinAlgError: | |
| raise SingularHessian() | |
| # Pyodide's NumPy can't see floating-point errors, so solve returns inf or nan instead of raising | |
| if np.all(np.isfinite(h)) and np.all(np.isfinite(g)) and not np.all(np.isfinite(delta)): | |
| raise SingularHessian() | |
| return p - delta | |
| class Adagrad: | |
| def __init__(self, objective: Objective, learning_rate: float, epsilon: float): | |
| self.f, self.lr, self.eps = objective, learning_rate, epsilon | |
| self.v = 0.0 # sum of squared gradients | |
| def step(self, p: np.ndarray) -> np.ndarray: | |
| g = self.f.gradient(p) | |
| self.v = self.v + g**2 | |
| return p - self.lr / np.sqrt(self.v + self.eps) * g | |
| class RMSProp: | |
| def __init__(self, objective: Objective, learning_rate: float, beta: float, epsilon: float): | |
| self.f, self.lr, self.beta, self.eps = objective, learning_rate, beta, epsilon | |
| self.v = 0.0 # moving average of squared gradients | |
| def step(self, p: np.ndarray) -> np.ndarray: | |
| g = self.f.gradient(p) | |
| self.v = self.beta * self.v + (1 - self.beta) * g**2 | |
| return p - self.lr / np.sqrt(self.v + self.eps) * g | |
| class Adadelta: | |
| def __init__(self, objective: Objective, beta: float, epsilon: float): | |
| self.f, self.beta, self.eps = objective, beta, epsilon | |
| self.v = 0.0 # moving average of squared gradients | |
| self.s = 0.0 # moving average of squared updates | |
| def step(self, p: np.ndarray) -> np.ndarray: | |
| g = self.f.gradient(p) | |
| self.v = self.beta * self.v + (1 - self.beta) * g**2 | |
| delta = -np.sqrt(self.s + self.eps) / np.sqrt(self.v + self.eps) * g | |
| self.s = self.beta * self.s + (1 - self.beta) * delta**2 | |
| return p + delta | |
| class Adam: | |
| def __init__(self, objective: Objective, learning_rate: float, beta1: float, beta2: float, epsilon: float): | |
| self.f, self.lr, self.beta1, self.beta2, self.eps = objective, learning_rate, beta1, beta2, epsilon | |
| self.m = 0.0 # first moment | |
| self.v = 0.0 # second moment | |
| self.t = 0 | |
| def step(self, p: np.ndarray) -> np.ndarray: | |
| g = self.f.gradient(p) | |
| self.t += 1 | |
| self.m = self.beta1 * self.m + (1 - self.beta1) * g | |
| self.v = self.beta2 * self.v + (1 - self.beta2) * g**2 | |
| m_hat = self.m / (1 - self.beta1**self.t) | |
| v_hat = self.v / (1 - self.beta2**self.t) | |
| return p - self.lr / (np.sqrt(v_hat) + self.eps) * m_hat | |
| # Ids match backend/src/algorithms.json | |
| OPTIMISERS = { | |
| "gradient_descent": GradientDescent, | |
| "nesterov": Nesterov, | |
| "newton": Newton, | |
| "adam": Adam, | |
| "adagrad": Adagrad, | |
| "rmsprop": RMSProp, | |
| "adadelta": Adadelta, | |
| } | |
| # A run that gets further than this from the origin has diverged | |
| DIVERGENCE_LIMIT = 1e6 | |
| DIVERGED = "The next point has no finite value or is beyond 1,000,000: the run diverged or left the function's domain." | |
| SINGULAR = "The Hessian is singular here, so Newton's method can't take a step." | |
| class Run: | |
| """One optimiser's trajectory, computed as far as it has been asked for.""" | |
| def __init__(self, objective: Objective, optimiser, start: np.ndarray): | |
| self.f = objective | |
| self.optimiser = optimiser | |
| self.points = [self._describe(start)] | |
| self.position = start | |
| # Why the run ended, once it has | |
| self.stopped: str | None = None | |
| def point(self, k: int) -> dict | None: | |
| """The point after k iterations, or None if the run stopped before then.""" | |
| with np.errstate(all="ignore"): | |
| while len(self.points) <= k and self.stopped is None: | |
| self._advance() | |
| return self.points[k] if k < len(self.points) else None | |
| def _advance(self) -> None: | |
| try: | |
| p = self.optimiser.step(self.position) | |
| except SingularHessian: | |
| self.stopped = SINGULAR | |
| return | |
| if not np.all(np.abs(p) <= DIVERGENCE_LIMIT) or not np.isfinite(self.f.value(p)): | |
| self.stopped = DIVERGED | |
| return | |
| self.position = p | |
| self.points.append(self._describe(p)) | |
| def _describe(self, p: np.ndarray) -> dict: | |
| return { | |
| "position": p.tolist(), | |
| "value": self.f.value(p), | |
| "gradient": self.f.gradient(p).tolist(), | |
| "hessian": self.f.hessian(p).tolist(), | |
| } | |