From 43e8f5fdb974889dcece640972df7325a62bb61f Mon Sep 17 00:00:00 2001
From: Ivan Razumov <iarspider@gmail.com>
Date: Thu, 23 Mar 2023 16:47:33 +0100
Subject: [PATCH] Add tests for (py3-)torch

---
 .../PythonAnalysis/test/BuildFile.xml         |  9 +++
 PhysicsTools/PythonAnalysis/test/testTorch.cc | 64 +++++++++++++++++++
 .../PythonAnalysis/test/test_torch.py         | 55 ++++++++++++++++
 3 files changed, 128 insertions(+)
 create mode 100644 PhysicsTools/PythonAnalysis/test/testTorch.cc
 create mode 100755 PhysicsTools/PythonAnalysis/test/test_torch.py

diff --git a/PhysicsTools/PythonAnalysis/test/BuildFile.xml b/PhysicsTools/PythonAnalysis/test/BuildFile.xml
index eea1fdc2970b1..0b2465e8227e5 100644
--- a/PhysicsTools/PythonAnalysis/test/BuildFile.xml
+++ b/PhysicsTools/PythonAnalysis/test/BuildFile.xml
@@ -14,6 +14,10 @@
 <test name="testTheano" command="testTheano.sh testTheano.py"/>
 <test name="testhep_ml" command="testTheano.sh testhep_ml.py"/>
 <test name="testCorrectionlib" command="test_correctionlib.py"/>
+<test name="testTorch" command="test_torch.py cpu"/>
+<iftool name="cuda">
+<test name="testTorchGPU" command="test_torch.py cuda:0"/>
+</iftool>
 
 <bin name="test_PyMVA" file="test_PyMVA.cpp">
   <use name="rootpymva"/>
@@ -139,6 +143,11 @@
 <test name="import-lxml" command="python3 -c 'import lxml'"/>
 <test name="import-bs4" command="python3 -c 'import bs4'"/>
 <test name="import-professor2" command="python3 -c 'import professor2'"/>
+<test name="import-torch" command="python3 -c 'import torch'"/>
 <test name="run-flawfinder" command="flawfinder -h"/>
 <test name="run-ipython" command="ipython -h"/>
 <test name="run-pylint" command="pylint -h"/>
+
+<bin name="testTorch" file="testTorch.cc">
+  <use name="torch"/>
+</bin>
diff --git a/PhysicsTools/PythonAnalysis/test/testTorch.cc b/PhysicsTools/PythonAnalysis/test/testTorch.cc
new file mode 100644
index 0000000000000..ae38f0f790393
--- /dev/null
+++ b/PhysicsTools/PythonAnalysis/test/testTorch.cc
@@ -0,0 +1,64 @@
+// Based on https://github.com/Maverobot/libtorch_examples/blob/master/src/simple_optimization_example.cpp
+#include <torch/torch.h>
+#include <cstdlib>
+#include <iostream>
+
+constexpr double kLearningRate = 0.001;
+constexpr int kMaxIterations = 100000;
+
+void native_run(double minimal) {
+  // Initial x value
+  auto x = torch::randn({1, 1}, torch::requires_grad(true));
+
+  for (size_t t = 0; t < kMaxIterations; t++) {
+    // Expression/value to be minimized
+    auto y = (x - minimal) * (x - minimal);
+    if (y.item<double>() < 1e-3) {
+      break;
+    }
+    // Calculate gradient
+    y.backward();
+
+    // Step x value without considering gradient
+    torch::NoGradGuard no_grad_guard;
+    x -= kLearningRate * x.grad();
+
+    // Reset the gradient of variable x
+    x.mutable_grad().reset();
+  }
+
+  std::cout << "[native] Actual minimal x value: " << minimal << ", calculated optimal x value: " << x.item<double>()
+            << std::endl;
+}
+
+void optimizer_run(double minimal) {
+  // Initial x value
+  std::vector<torch::Tensor> x;
+  x.push_back(torch::randn({1, 1}, torch::requires_grad(true)));
+  auto opt = torch::optim::SGD(x, torch::optim::SGDOptions(kLearningRate));
+
+  for (size_t t = 0; t < kMaxIterations; t++) {
+    // Expression/value to be minimized
+    auto y = (x[0] - minimal) * (x[0] - minimal);
+    if (y.item<double>() < 1e-3) {
+      break;
+    }
+    // Calculate gradient
+    y.backward();
+
+    // Step x value without considering gradient
+    opt.step();
+    // Reset the gradient of variable x
+    opt.zero_grad();
+  }
+
+  std::cout << "[optimizer] Actual minimal x value: " << minimal
+            << ", calculated optimal x value: " << x[0].item<double>() << std::endl;
+}
+
+// optimize y = (x - 10)^2
+int main(int argc, char* argv[]) {
+  native_run(0.01);
+  optimizer_run(0.01);
+  return 0;
+}
diff --git a/PhysicsTools/PythonAnalysis/test/test_torch.py b/PhysicsTools/PythonAnalysis/test/test_torch.py
new file mode 100755
index 0000000000000..8726982e2c276
--- /dev/null
+++ b/PhysicsTools/PythonAnalysis/test/test_torch.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+
+import torch
+import math
+import sys
+
+
+dtype = torch.float
+device = torch.device(sys.argv[1])
+
+# Create random input and output data
+try:
+    x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
+except RuntimeError as e:
+    if 'cuda' in str(e).lower():
+        print("CUDA-related error - is CUDA device available?")
+        print(str(e))
+        exit(0)
+    else:
+        raise e
+
+y = torch.sin(x)
+
+# Randomly initialize weights
+a = torch.randn((), device=device, dtype=dtype)
+b = torch.randn((), device=device, dtype=dtype)
+c = torch.randn((), device=device, dtype=dtype)
+d = torch.randn((), device=device, dtype=dtype)
+
+learning_rate = 1e-6
+for t in range(2000):
+    # Forward pass: compute predicted y
+    y_pred = a + b * x + c * x ** 2 + d * x ** 3
+
+    # Compute and print loss
+    loss = (y_pred - y).pow(2).sum().item()
+    if t % 100 == 99:
+        print(t, loss)
+
+    # Backprop to compute gradients of a, b, c, d with respect to loss
+    grad_y_pred = 2.0 * (y_pred - y)
+    grad_a = grad_y_pred.sum()
+    grad_b = (grad_y_pred * x).sum()
+    grad_c = (grad_y_pred * x ** 2).sum()
+    grad_d = (grad_y_pred * x ** 3).sum()
+
+    # Update weights using gradient descent
+    a -= learning_rate * grad_a
+    b -= learning_rate * grad_b
+    c -= learning_rate * grad_c
+    d -= learning_rate * grad_d
+
+
+print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')
+