Merge 36bfd5d into f12b175

mlcalcdriver/base/job.py

@@ -133,7 +133,7 @@ def calculator(self, calculator):
                 """
             )
 
-    def run(self, property, device="cpu", batch_size=128, finite_difference=False):
+    def run(self, property, batch_size=128, finite_difference=False):
         r"""
         Main method to call to obtain results for a Job
 
@@ -144,20 +144,10 @@ def run(self, property, device="cpu", batch_size=128, finite_difference=False):
             available_properties of the Calculator except the
             forces which can be derived from an energy
             Calculator.
-        device : str
-            Device on which to run the calculation.
-            Either `"cpu"` or `"cuda"` to run on cpu or gpu.
-            Default is `"cpu"` and should not be changed, except
-            for very large systems.
         batch_size : int
             Size of the mini-batches used in predictions.
             Default is 128.
         """
-        device = str(device)
-        if device.startswith("cuda"):
-            if not torch.cuda.is_available():
-                warnings.warn("CUDA was asked for, but is not available.", UserWarning)
-
         if not finite_difference:
             predictions = self.calculator.run(
                 property=property, posinp=self.posinp, batch_size=batch_size

mlcalcdriver/calculators/schnetpack.py

@@ -35,17 +35,26 @@ def __init__(self, model_dir, available_properties=None, device="cpu", units=eVA
         device : str
             Can be either `"cpu"` to use cpu or `"cuda"` to use "gpu"
         """
+        self.device = device
         try:
-            self.model = load_model(model_dir=model_dir, device=device)
+            self.model = load_model(model_dir=model_dir, device=self.device)
         except Exception:
             self.model = load_model(
-                model_dir=os.environ["MODELDIR"] + model_dir, device=device
+                model_dir=os.environ["MODELDIR"] + model_dir, device=self.device
             )
         super(SchnetPackCalculator, self).__init__(units=units)
         self._get_representation_type()
 
+    @property
+    def device(self):
+        return self._device
+
+    @device.setter
+    def device(self, device):
+        self._device = str(device).lower()
+
     def run(
-        self, property, posinp=None, device="cpu", batch_size=128,
+        self, property, posinp=None, batch_size=128,
     ):
         r"""
         Main method to use when making a calculation with
@@ -76,16 +85,16 @@ def run(
         if derivative == 0:
             if self.model.output_modules[0].derivative is not None:
                 for batch in data_loader:
-                    batch = {k: v.to(device) for k, v in batch.items()}
+                    batch = {k: v.to(self.device) for k, v in batch.items()}
                     pred.append(self.model(batch))
             else:
                 with torch.no_grad():
                     for batch in data_loader:
-                        batch = {k: v.to(device) for k, v in batch.items()}
+                        batch = {k: v.to(self.device) for k, v in batch.items()}
                         pred.append(self.model(batch))
         if abs(derivative) == 1:
             for batch in data_loader:
-                batch = {k: v.to(device) for k, v in batch.items()}
+                batch = {k: v.to(self.device) for k, v in batch.items()}
                 batch[wrt[0]].requires_grad_()
                 results = self.model(batch)
                 deriv1 = torch.unsqueeze(
@@ -96,7 +105,7 @@ def run(
                 pred.append({out_name: deriv1})
         if abs(derivative) == 2:
             for batch in data_loader:
-                batch = {k: v.to(device) for k, v in batch.items()}
+                batch = {k: v.to(self.device) for k, v in batch.items()}
                 for inp in set(wrt):
                     batch[inp].requires_grad_()
                 results = self.model(batch)

mlcalcdriver/workflows/geopt.py

@@ -140,13 +140,10 @@ def max_iter(self):
     def max_iter(self, max_iter):
         self._max_iter = int(max_iter)
 
-    def run(self, device="cpu", batch_size=128, recenter=False, verbose=0):
+    def run(self, batch_size=128, recenter=False, verbose=0):
         r"""
         Parameters
         ----------
-        device : str
-            Either 'cpu' or 'cuda' to run on cpu or gpu. Default is 'cpu'
-            and should be faster in most cases.
         batch_size : int
             Size of the mini-batches used in predictions. Default is 128.
         recenter : bool
@@ -166,7 +163,7 @@ def run(self, device="cpu", batch_size=128, recenter=False, verbose=0):
         for i in range(1, self.max_iter + 1):
             # Forces calculation
             job = Job(posinp=temp_posinp, calculator=self.calculator)
-            job.run("forces", device=device, batch_size=batch_size)
+            job.run("forces", batch_size=batch_size)
             # Moving the atoms
             for j in range(len(job.posinp[0])):
                 temp_posinp = temp_posinp.translate_atom(

mlcalcdriver/workflows/phonon.py

@@ -219,13 +219,10 @@ def normal_modes(self):
     def normal_modes(self, normal_modes):
         self._normal_modes = normal_modes
 
-    def run(self, device="cpu", batch_size=128, **kwargs):
+    def run(self, batch_size=128, **kwargs):
         r"""
         Parameters
         ----------
-        device : str
-            Either "cpu" or "cuda" to run on cpu or gpu. Default is 'cpu'
-            and should be faster in most cases.
         batch_size : int
             Batch size used when passing the structures to the model
         **kwargs :
@@ -234,15 +231,15 @@ def run(self, device="cpu", batch_size=128, **kwargs):
         """
         if self.relax:
             geopt = Geopt(posinp=self.posinp, calculator=self.calculator, **kwargs)
-            geopt.run(device=device, batch_size=batch_size)
+            geopt.run(batch_size=batch_size)
             self._ground_state = deepcopy(geopt.final_posinp)
 
         if self.finite_difference:
             job = Job(posinp=self._create_displacements(), calculator=self.calculator)
-            job.run(property="forces", device=device, batch_size=batch_size)
+            job.run(property="forces", batch_size=batch_size)
         else:
             job = Job(posinp=self._ground_state, calculator=self.calculator)
-            job.run(property="hessian", device=device, batch_size=batch_size)
+            job.run(property="hessian", batch_size=batch_size)
         self._post_proc(job)
 
     def _create_displacements(self):

requirements_dev.txt

@@ -1,15 +1,15 @@
 pip==20.1.1
 wheel==0.34.2
 black==19.10b0
-numpy==1.18.5
-torch==1.5.0
-torchvision==0.6.0
-twine==3.1.1
+numpy==1.19.0
+torch==1.5.1
+torchvision==0.6.1
+twine==3.2.0
 pytest==5.4.3
 pytest-cov==2.10.0
-pytest-sugar==0.9.3
-coveralls==2.0.0
+pytest-sugar==0.9.4
+coveralls==2.1.1
 ase==3.19.1
 schnetpack==0.3
-sphinx==3.1.1
-sphinx-rtd-theme==0.4.3
+sphinx==3.1.2
+sphinx-rtd-theme==0.5.0

tests/test_schnetpack.py

@@ -24,10 +24,6 @@ def test_only_energy(self):
         self.job.run("energy")
         assert np.array(-2979.6067) == self.job.results["energy"][0]
 
-    def test_cuda_unavailable(self):
-        with pytest.warns(UserWarning):
-            self.job.run("energy", device="cuda")
-
     def test_forces_from_deriv(self):
         assert self.job.calculator.available_properties == ["energy"]
         self.job.run("forces")