class_docstr: This is the base class for the exchange-correlation (XC) functional.
  The XC functional is used to calculate the exchange-correlation energy
  and potential. The XC functional is usually divided into three families:
  LDA, GGA, and Meta-GGA. The LDA is the simplest one, which only depends
  on the density. The GGA depends on the density and its gradient. The
  Meta-GGA depends on the density, its gradient, and its Laplacian.
  
  Examples
  --------
  >>> import torch
  >>> from deepchem.utils.dft_utils import ValGrad, SpinParam
  >>> from deepchem.utils.dft_utils import BaseXC
  >>> class MyXC(BaseXC):
  ...     @property
  ...     def family(self) -> int:
  ...         return 1
  ...     def get_edensityxc(self, densinfo: Union[ValGrad, SpinParam[ValGrad]]) -> torch.Tensor:
  ...         if isinstance(densinfo, ValGrad):
  ...             return densinfo.value.pow(2)
  ...         else:
  ...             return densinfo.u.value.pow(2) + densinfo.d.value.pow(2)
  ...     def get_vxc(self, densinfo: Union[ValGrad, SpinParam[ValGrad]]) -> Union[ValGrad, SpinParam[ValGrad]]:
  ...         if isinstance(densinfo, ValGrad):
  ...             return ValGrad(value=2*densinfo.value)
  ...         else:
  ...             return SpinParam(u=ValGrad(value=2*densinfo.u.value),
  ...                              d=ValGrad(value=2*densinfo.d.value))
  >>> xc = MyXC()
  >>> densinfo = ValGrad(value=torch.tensor([1., 2., 3.], requires_grad=True))
  >>> xc.get_edensityxc(densinfo)
  tensor([1., 4., 9.], grad_fn=<PowBackward0>)
  >>> xc.get_vxc(densinfo)
  ValGrad(value=tensor([2., 4., 6.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None)
  >>> densinfo = SpinParam(u=ValGrad(value=torch.tensor([1., 2., 3.], requires_grad=True)),
  ...                      d=ValGrad(value=torch.tensor([4., 5., 6.], requires_grad=True)))
  >>> xc.get_edensityxc(densinfo)
  tensor([17., 29., 45.], grad_fn=<AddBackward0>)
  >>> xc.get_vxc(densinfo)
  SpinParam(u=ValGrad(value=tensor([2., 4., 6.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None), d=ValGrad(value=tensor([ 8., 10., 12.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None))
class_docstr_tok: ["This","is","the","base","class","for","the","exchange-correlation","(","XC",")","functional",".","The","XC","functional","is","used","to","calculate","the","exchange-correlation","energy","and","potential",".","The","XC","functional","is","usually","divided","into","three","families",":","LDA",",","GGA",",","and","Meta-GGA",".","The","LDA","is","the","simplest","one",",","which","only","depends","on","the","density",".","The","GGA","depends","on","the","density","and","its","gradient",".","The","Meta-GGA","depends","on","the","density",",","its","gradient",",","and","its","Laplacian",".","Examples","--","--","--","--",">",">",">","import","torch",">",">",">","from","deepchem.utils.dft_utils","import","ValGrad",",","SpinParam",">",">",">","from","deepchem.utils.dft_utils","import","BaseXC",">",">",">","class","MyXC","(","BaseXC",")",":","...","@","property","...","def","family","(","self",")","-",">","int",":","...","return","1","...","def","get_edensityxc","(","self",",","densinfo",":","Union","[","ValGrad",",","SpinParam","[","ValGrad","]","]",")","-",">","torch.Tensor",":","...","if","isinstance","(","densinfo",",","ValGrad",")",":","...","return","densinfo.value.pow","(","2",")","...","else",":","...","return","densinfo.u.value.pow","(","2",")","+","densinfo.d.value.pow","(","2",")","...","def","get_vxc","(","self",",","densinfo",":","Union","[","ValGrad",",","SpinParam","[","ValGrad","]","]",")","-",">","Union","[","ValGrad",",","SpinParam","[","ValGrad","]","]",":","...","if","isinstance","(","densinfo",",","ValGrad",")",":","...","return","ValGrad","(","value=2","*","densinfo.value",")","...","else",":","...","return","SpinParam","(","u=ValGrad","(","value=2","*","densinfo.u.value",")",",","...","d=ValGrad","(","value=2","*","densinfo.d.value",")",")",">",">",">","xc","=","MyXC","(",")",">",">",">","densinfo","=","ValGrad","(","value=torch.tensor","(","[","1.",",","2.",",","3",".","]",",","requires_grad=True",")",")",">",">",">","xc.get_edensityxc","(","densinfo",")","tensor","(","[","1.",",","4.",",","9",".","]",",","grad_fn=","<","PowBackward0",">",")",">",">",">","xc.get_vxc","(","densinfo",")","ValGrad","(","value=tensor","(","[","2.",",","4.",",","6",".","]",",","grad_fn=","<","MulBackward0",">",")",",","grad=None",",","lapl=None",",","kin=None",")",">",">",">","densinfo","=","SpinParam","(","u=ValGrad","(","value=torch.tensor","(","[","1.",",","2.",",","3",".","]",",","requires_grad=True",")",")",",","...","d=ValGrad","(","value=torch.tensor","(","[","4.",",","5.",",","6",".","]",",","requires_grad=True",")",")",")",">",">",">","xc.get_edensityxc","(","densinfo",")","tensor","(","[","17.",",","29.",",","45",".","]",",","grad_fn=","<","AddBackward0",">",")",">",">",">","xc.get_vxc","(","densinfo",")","SpinParam","(","u=ValGrad","(","value=tensor","(","[","2.",",","4.",",","6",".","]",",","grad_fn=","<","MulBackward0",">",")",",","grad=None",",","lapl=None",",","kin=None",")",",","d=ValGrad","(","value=tensor","(","[","8.",",","10.",",","12",".","]",",","grad_fn=","<","MulBackward0",">",")",",","grad=None",",","lapl=None",",","kin=None",")",")"]
docstr: set the context where some elements (depends on xc family) in densinfo requires grad
  
  Parameters
  ----------
  densinfo : Union[ValGrad, SpinParam[ValGrad]]
      The density information.
      If the XC is unpolarized, then densinfo is ValGrad.
      If the XC is polarized, then densinfo is SpinParam[ValGrad].
  
  Returns
  -------
  Iterator
      The context where some elements in densinfo requires grad
docstr_tok: ["set","the","context","where","some","elements","(","depends","on","xc","family",")","in","densinfo","requires","grad","Parameters","--","--","--","--","--","densinfo",":","Union","[","ValGrad",",","SpinParam","[","ValGrad","]","]","The","density","information",".","If","the","XC","is","unpolarized",",","then","densinfo","is","ValGrad",".","If","the","XC","is","polarized",",","then","densinfo","is","SpinParam","[","ValGrad","]",".","Returns","--","--","--","-","Iterator","The","context","where","some","elements","in","densinfo","requires","grad"]
code: def _enable_grad_densinfo(
      self, densinfo: Union[ValGrad, SpinParam[ValGrad]]
  ) -> Iterator:
      """
      set the context where some elements (depends on xc family) in densinfo requires grad
  
      Parameters
      ----------
      densinfo : Union[ValGrad, SpinParam[ValGrad]]
          The density information.
          If the XC is unpolarized, then densinfo is ValGrad.
          If the XC is polarized, then densinfo is SpinParam[ValGrad].
  
      Returns
      -------
      Iterator
          The context where some elements in densinfo requires grad
  
      """
  
      def _get_set_grad(vars: List[torch.Tensor]) -> List[bool]:
          # set the vars to require grad and returns the previous state of the vars
          reqgrads = []
          for var in vars:
              reqgrads.append(var.requires_grad)
              var.requires_grad_()
          return reqgrads
  
      def _restore_grad(
          reqgrads: List[bool], vars: List[torch.Tensor]
      ) -> None:
          # restore the state of requiring grad based on reqgrads list
          # all vars before this function requires grad
          for reqgrad, var in zip(reqgrads, vars):
              if not reqgrad:
                  var.requires_grad_(False)
  
      # getting which parameters should require grad
      if not isinstance(densinfo, ValGrad):  # a spinparam
          params = [densinfo.u.value, densinfo.d.value]
          if self.family >= 2:  # GGA
              assert densinfo.u.grad is not None
              assert densinfo.d.grad is not None
              params.extend([densinfo.u.grad, densinfo.d.grad])
          if self.family >= 3:  # MGGA
              assert densinfo.u.lapl is not None
              assert densinfo.d.lapl is not None
              assert densinfo.u.kin is not None
              assert densinfo.d.kin is not None
              params.extend(
                  [
                      densinfo.u.lapl,
                      densinfo.d.lapl,
                      densinfo.u.kin,
                      densinfo.d.kin,
                  ]
              )
      else:
          params = [densinfo.value]
          if self.family >= 2:
              assert densinfo.grad is not None
              params.append(densinfo.grad)
          if self.family >= 3:
              assert densinfo.lapl is not None
              assert densinfo.kin is not None
              params.extend([densinfo.lapl, densinfo.kin])
  
      try:
          # set the params to require grad
          reqgrads = _get_set_grad(params)
          yield
      finally:
          _restore_grad(reqgrads, params)
  
imports: from contextlib import contextmanager
  from abc import abstractmethod, abstractproperty
  import torch
  from deepchem.utils.differentiation_utils import EditableModule
  from typing import Any, List, Union, Iterator
  from deepchem.utils.dft_utils import ValGrad, SpinParam
label_desc: Use image node_id 5 for calling the BaseXC obj's underlying member method code with example usage: obj._enable_grad_densinfo(densinfo) and returns: reqgrads