warp-lang 0.15.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/sim/optimizer.py

@@ -1,138 +0,0 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-from warp.context import synchronize
-import warp as wp
-
-import numpy as np
-
-
-@wp.kernel
-def gd_step(arr_x: wp.array(dtype=float), arr_dfdx: wp.array(dtype=float), alpha: float):
-    tid = wp.tid()
-
-    x = arr_x[tid]
-    dfdx = arr_dfdx[tid]
-
-    x = x - dfdx * alpha
-
-    arr_x[tid] = x
-
-
-@wp.kernel
-def nesterov1(beta: float, x: wp.array(dtype=float), x_prev: wp.array(dtype=float), y: wp.array(dtype=float)):
-    tid = wp.tid()
-
-    y[tid] = x[tid] + beta * (x[tid] - x_prev[tid])
-
-
-@wp.kernel
-def nesterov2(
-    alpha: float,
-    beta: wp.array(dtype=float),
-    eta: wp.array(dtype=float),
-    x: wp.array(dtype=float),
-    x_prev: wp.array(dtype=float),
-    y: wp.array(dtype=float),
-    dfdx: wp.array(dtype=float),
-):
-    # if (eta > 0.0):
-    #     # adaptive restart
-    #     x_prev = x
-    #     b = 0
-    # else:
-    #     # nesterov update
-    #     x_prev = x
-    #     x = y - alpha*dfdx
-
-    tid = wp.tid()
-
-    x_prev[tid] = x[tid]
-    x[tid] = y[tid] - alpha * dfdx[tid]
-
-
-def inner(a, b, out):
-    from warp.utils import array_inner
-
-    array_inner(a, b, out)
-
-
-class Optimizer:
-    def __init__(self, n, mode, device):
-        self.n = n
-        self.mode = mode
-        self.device = device
-
-        # allocate space for residual buffers
-        self.dfdx = wp.zeros(n, dtype=float, device=device)
-
-        if mode == "nesterov":
-            self.x_prev = wp.zeros(n, dtype=float, device=device)
-            self.y = wp.zeros(n, dtype=float, device=device)
-            self.eta = wp.zeros(1, dtype=float, device=device)
-            self.eta_prev = wp.zeros(1, dtype=float, device=device)
-            self.beta = wp.zeros(1, dtype=int, device=device)
-
-    def solve(self, x, grad_func, max_iters=20, alpha=0.01, report=False):
-        if report:
-            stats = {}
-
-            # reset stats
-            stats["evals"] = 0
-            stats["residual"] = []
-
-        if self.mode == "gd":
-            for i in range(max_iters):
-                # compute residual
-                grad_func(x, self.dfdx)
-
-                # gradient step
-                wp.launch(kernel=gd_step, dim=self.n, inputs=[x, self.dfdx, alpha], device=self.device)
-
-                if report:
-                    stats["evals"] += 1
-
-                    r = np.linalg.norm(self.dfdx.to("cpu").numpy())
-                    stats["residual"].append(r)
-
-        elif self.mode == "nesterov":
-            wp.copy(self.x_prev, x)
-
-            # momentum index (reset after restart)
-            b = 0
-            for iter in range(max_iters):
-                beta = (b - 1.0) / (b + 2.0)
-                b += 1
-
-                # y = x + beta*(x - x_prev)
-                wp.launch(kernel=nesterov1, dim=self.n, inputs=[beta, x, self.x_prev, self.y], device=self.device)
-
-                # grad
-                grad_func(self.y, self.dfdx)
-
-                # inner()
-                # np.dot(dfdx, x - x_prev)
-
-                # x = y - alpha*dfdx
-                wp.launch(
-                    kernel=nesterov2,
-                    dim=self.n,
-                    inputs=[alpha, None, None, x, self.x_prev, self.y, self.dfdx],
-                    device=self.device,
-                )
-
-                if report:
-                    stats["evals"] += 1
-
-                    r = np.linalg.norm(self.dfdx.to("cpu").numpy())
-                    stats["residual"].append(r)
-
-        else:
-            raise RuntimeError("Unknown optimizer")
-
-        if report:
-            print(stats)