Cleanup-1 [termination condition and init file] (deepchem#4109)

* minimizer and init file

* added test

* added test

* fixed formatting

deepchem/utils/differentiation_utils/optimize/minimizer.py

@@ -18,6 +18,7 @@ def gd(
         x_rtol: float = 1e-8,
         # misc parameters
         verbose=False,
+        terminate=False,
         **unused):
     r"""
     Vanilla gradient descent with momentum. The stopping conditions use OR criteria.
@@ -61,6 +62,8 @@ def gd(
         The absolute tolerance of the norm of the input ``x``.
     x_rtol: float or None
         The relative tolerance of the norm of the input ``x``.
+    terminate: bool (default False)
+        Whether to use termination condition or, keep on running the minimizer.
 
     """
 
@@ -79,11 +82,12 @@ def gd(
         # check the stopping conditions
         to_stop = stop_cond.to_stop(i, x, xprev, f, fprev)
 
-        if to_stop:
+        if to_stop and terminate:
             break
 
         fprev = f
-    x = stop_cond.get_best_x(x)
+    if terminate:
+        x = stop_cond.get_best_x(x)
     return x
 
 
@@ -104,6 +108,7 @@ def adam(
         x_rtol: float = 1e-8,
         # misc parameters
         verbose=False,
+        terminate=False,
         **unused):
     r"""
     Adam optimizer by Kingma & Ba (2015). The stopping conditions use OR criteria.
@@ -148,6 +153,8 @@ def adam(
         The absolute tolerance of the norm of the input ``x``.
     x_rtol: float or None
         The relative tolerance of the norm of the input ``x``.
+    terminate: bool (default False)
+        Whether to use the termination condition, or keep running the minimizer.
     """
 
     x = x0.clone()
@@ -175,11 +182,12 @@ def adam(
         # check the stopping conditions
         to_stop = stop_cond.to_stop(i, x, xprev, f, fprev)
 
-        if to_stop:
+        if to_stop and terminate:
             break
 
         fprev = f
-    x = stop_cond.get_best_x(x)
+    if terminate:
+        x = stop_cond.get_best_x(x)
     return x
 
 

deepchem/utils/test/test_differentiation_utils.py

@@ -1228,3 +1228,18 @@ def lotka_volterra(t, z, *params):
     assert torch.allclose(sol[-1][0],
                           torch.tensor(sol_scipy.y[0][-1], dtype=torch.float),
                           0.01, 0.001)
+
+
+@pytest.mark.torch
+def test_terminate_param():
+    from deepchem.utils.differentiation_utils import gd
+    import torch
+
+    def fun(x):
+        return torch.tan(x), (1 / torch.cos(x))**2
+
+    x0 = torch.tensor(0.0, requires_grad=True)
+    x0.grad = torch.tensor(1.0)
+    x1 = gd(fun, x0, [], terminate=True)
+    x2 = gd(fun, x0, [], terminate=False)
+    assert not torch.allclose(x1, x2)