class_docstr: A deep-learning based Variational Monte Carlo method [1]_ for calculating the ab-initio
  solution of a many-electron system.
  
  This model aims to calculate the ground state energy of a multi-electron system
  using a baseline solution as the Hartree-Fock. An MCMC technique is used to sample
  electrons and DNNs are used to caluclate the square magnitude of the wavefunction,
  in which electron-electron repulsions also are included in the calculation(in the
  form of Jastrow factor envelopes). The model requires only the nucleus' coordinates
  as input.
  
  This method is based on the following paper:
  
  Example
  -------
  >>> from deepchem.models.torch_models.ferminet import FerminetModel
  >>> H2_molecule = [['H', [0, 0, 0]], ['H', [0, 0, 0.748]]]
  >>> mol = FerminetModel(H2_molecule, spin=0, ion_charge=0, tasks='pretraining')
  >>> mol.train(nb_epoch=3)
  >>> print(mol.model.psi_up.size())
  torch.Size([1, 1])
  
  References
  ----------
  .. [1] Spencer, James S., et al. Better, Faster Fermionic Neural Networks. arXiv:2011.07125, arXiv, 13 Nov. 2020. arXiv.org, http://arxiv.org/abs/2011.07125.
  
  Note
  ----
  This class requires pySCF to be installed.
class_docstr_tok: ["A","deep-learning","based","Variational","Monte","Carlo","method","[","1","]","_","for","calculating","the","ab-initio","solution","of","a","many-electron","system",".","This","model","aims","to","calculate","the","ground","state","energy","of","a","multi-electron","system","using","a","baseline","solution","as","the","Hartree-Fock",".","An","MCMC","technique","is","used","to","sample","electrons","and","DNNs","are","used","to","caluclate","the","square","magnitude","of","the","wavefunction",",","in","which","electron-electron","repulsions","also","are","included","in","the","calculation","(","in","the","form","of","Jastrow","factor","envelopes",")",".","The","model","requires","only","the","nucleus","'","coordinates","as","input",".","This","method","is","based","on","the","following","paper",":","Example","--","--","--","-",">",">",">","from","deepchem.models.torch_models.ferminet","import","FerminetModel",">",">",">","H2_molecule","=","[","[","'","H","'",",","[","0",",","0",",","0","]","]",",","[","'","H","'",",","[","0",",","0",",","0.748","]","]","]",">",">",">","mol","=","FerminetModel","(","H2_molecule",",","spin=0",",","ion_charge=0",",","tasks='pretraining","'",")",">",">",">","mol.train","(","nb_epoch=3",")",">",">",">","print","(","mol.model.psi_up.size","(",")",")","torch.Size","(","[","1",",","1","]",")","References","--","--","--","--","--","..","[","1","]","Spencer",",","James","S.",",","et","al",".","Better",",","Faster","Fermionic","Neural","Networks",".","arXiv:2011.07125",",","arXiv",",","13","Nov.","2020.","arXiv.org",",","http",":","\/\/arxiv.org\/abs\/2011.07125",".","Note","--","--","This","class","requires","pySCF","to","be","installed","."]
docstr: function to run training or pretraining.
  
  Parameters
  ----------
  nb_epoch: int (default: 200)
      contains the number of pretraining steps to be performed
  lr : float (default: 0.002)
      contains the learning rate for the model fitting
  weight_decay: float (default: 0.002)
      contains the weight_decay for the model fitting
  std: float (default: 0.08)
      The standard deviation for the electron update during training
  std_init: float (default: 0.02)
      The standard deviation for the electron update during pretraining
  steps_std: float (default 100)
      The number of steps for standard deviation increase
docstr_tok: ["function","to","run","training","or","pretraining",".","Parameters","--","--","--","--","--","nb_epoch",":","int","(","default",":","200",")","contains","the","number","of","pretraining","steps","to","be","performed","lr",":","float","(","default",":","0.002",")","contains","the","learning","rate","for","the","model","fitting","weight_decay",":","float","(","default",":","0.002",")","contains","the","weight_decay","for","the","model","fitting","std",":","float","(","default",":","0.08",")","The","standard","deviation","for","the","electron","update","during","training","std_init",":","float","(","default",":","0.02",")","The","standard","deviation","for","the","electron","update","during","pretraining","steps_std",":","float","(","default","100",")","The","number","of","steps","for","standard","deviation","increase"]
code: def train(
      self,
      nb_epoch: int = 200,
      lr: float = 0.002,
      weight_decay: float = 0,
      std: float = 0.08,
      std_init: float = 0.02,
      steps_std: int = 100,
  ):
      """
      function to run training or pretraining.
  
      Parameters
      ----------
      nb_epoch: int (default: 200)
          contains the number of pretraining steps to be performed
      lr : float (default: 0.002)
          contains the learning rate for the model fitting
      weight_decay: float (default: 0.002)
          contains the weight_decay for the model fitting
      std: float (default: 0.08)
          The standard deviation for the electron update during training
      std_init: float (default: 0.02)
          The standard deviation for the electron update during pretraining
      steps_std: float (default 100)
          The number of steps for standard deviation increase
      """
  
      # hook function below is an efficient way modifying the gradients on the go rather than looping
      def energy_hook(grad, random_walk_steps):
          """
          hook function to modify the gradients
          """
          # using non-local variables as a means of parameter passing
          nonlocal energy_local, energy_mean
          new_grad = (2 / random_walk_steps) * (
              (energy_local - energy_mean) * grad
          )
          return new_grad.float()
  
      optimizer = torch.optim.Adam(
          self.model.parameters(), lr=lr, weight_decay=weight_decay
      )
  
      if self.tasks == "pretraining":
          for iteration in range(nb_epoch):
              optimizer.zero_grad()
              accept = self.molecule.move(stddev=std_init)
              if iteration % steps_std == 0:
                  if accept > 0.55:
                      std_init *= 1.1
                  else:
                      std_init /= 1.1
              self.loss_value = (
                  torch.mean(self.model.running_diff)
                  / self.random_walk_steps
              )
              self.loss_value.backward()
              optimizer.step()
              logging.info(
                  "The loss for the pretraining iteration "
                  + str(iteration)
                  + " is "
                  + str(self.loss_value.item())
              )
              self.model.running_diff = torch.zeros(self.batch_no)
  
      if self.tasks == "training":
          energy_local = None
          optimizer = torch.optim.Adam(
              self.model.parameters(), lr=lr, weight_decay=weight_decay
          )
          self.final_energy = torch.tensor(0.0)
          with torch.no_grad():
              hooks = list(
                  map(
                      lambda param: param.register_hook(
                          lambda grad: energy_hook(
                              grad, self.random_walk_steps
                          )
                      ),
                      self.model.parameters(),
                  )
              )
          for iteration in range(nb_epoch):
              optimizer.zero_grad()
              self.energy_sampled = torch.tensor([])
              # the move function calculates the energy of sampled electrons and samples new set of electrons (does not calculate loss)
              accept = self.molecule.move(stddev=std)
              if iteration % steps_std == 0:
                  if accept > 0.55:
                      std_init *= 1.2
                  else:
                      std_init /= 1.2
              median, _ = torch.median(self.energy_sampled, 0)
              variance = torch.mean(
                  torch.abs(self.energy_sampled - median)
              )
              # clipping local energies which are away 5 times the variance from the median
              clamped_energy = torch.clamp(
                  self.energy_sampled,
                  max=median + 5 * variance,
                  min=median - 5 * variance,
              )
              energy_mean = torch.mean(clamped_energy)
              logging.info(
                  "The mean energy for the training iteration "
                  + str(iteration)
                  + " is "
                  + str(energy_mean.item())
              )
              self.final_energy = self.final_energy + energy_mean
              # using the sampled electrons from the electron sampler for bacward pass and modifying gradients
              sample_history = torch.from_numpy(
                  self.molecule.sampled_electrons
              ).view(self.random_walk_steps, self.batch_no, -1, 3)
              optimizer.zero_grad()
              for i in range(self.random_walk_steps):
                  # going through each step of random walk and calculating the modified gradients with local energies
                  input_electron = sample_history[i]
                  input_electron.requires_grad = True
                  energy_local = torch.mean(clamped_energy[i])
                  self.model.forward(input_electron)
                  self.loss_value = torch.mean(
                      torch.log(torch.abs(self.model.psi))
                  )
                  self.loss_value.backward()
              optimizer.step()
          self.final_energy = self.final_energy / nb_epoch
          list(map(lambda hook: hook.remove(), hooks))
  
imports: import logging
  from typing import List, Optional, Tuple
  from rdkit import Chem
  import numpy
  from deepchem.utils.molecule_feature_utils import ALLEN_ELECTRONEGATIVTY
  from deepchem.models.torch_models import TorchModel
  import torch
  from deepchem.models.torch_models.layers import FerminetElectronFeature, FerminetEnvelope
  from deepchem.utils.electron_sampler import ElectronSampler
label_desc: Use image node_id 6 for calling the FerminetModel obj's underlying member method code with example usage: obj.train(nb_epoch, lr, weight_decay, std, std_init, steps_std) and returns: new_grad