def _core_wavefunction_set_variable( cls: core.Wavefunction, key: str, val: Union[core.Matrix, np.ndarray, float]) -> None: """Sets scalar or array QCVariable *key* to *val* on *cls*.""" if isinstance(val, core.Matrix): if cls.has_scalar_variable(key): raise ValidationError( "psi4.core.Wavefunction.set_variable: Target variable " + key + " already a scalar variable!") else: cls.set_array_variable(key, val) elif isinstance(val, np.ndarray): if cls.has_scalar_variable(key): raise ValidationError( "psi4.core.Wavefunction.set_variable: Target variable " + key + " already a scalar variable!") else: cls.set_array_variable( key, core.Matrix.from_array(_qcvar_reshape_set(key, val))) else: if cls.has_array_variable(key): raise ValidationError( "psi4.core.Wavefunction.set_variable: Target variable " + key + " already an array variable!") else: cls.set_scalar_variable(key, val)
def _core_wavefunction_variable( cls: core.Wavefunction, key: str) -> Union[float, core.Matrix, np.ndarray]: """Return copy of scalar or array QCVariable *key* from *self* :class:`psi4.core.Wavefunction`. Returns ------- float or numpy.ndarray or Matrix Scalar variables are returned as floats. Array variables not naturally 2D (like multipoles) are returned as :class:`numpy.ndarray` of natural dimensionality. Other array variables are returned as :py:class:`~psi4.core.Matrix` and may have an extra dimension with symmetry information. Example ------- >>> g, wfn = psi4.gradient("hf/cc-pvdz", return_wfn=True) >>> wfn.variable("CURRENT ENERGY") -100.00985995185668 >>> wfn.variable("CURRENT DIPOLE") array([ 0. , 0. , -0.83217802]) >>> wfn.variable("CURRENT GRADIENT") <psi4.core.Matrix object at 0x12d884fc0> >>> wfn.variable("CURRENT GRADIENT").np array([[ 6.16297582e-33, 6.16297582e-33, -9.41037138e-02], [-6.16297582e-33, -6.16297582e-33, 9.41037138e-02]]) """ key = _qcvar_warnings(key) if cls.has_scalar_variable(key): return cls.scalar_variable(key) elif cls.has_array_variable(key): return _qcvar_reshape_get(key, cls.array_variable(key)) else: raise KeyError("psi4.core.Wavefunction.variable: Requested variable " + key + " was not set!\n")
def _core_wavefunction_del_variable(cls: core.Wavefunction, key: str) -> None: """Removes scalar or array QCVariable *key* from *cls* if present.""" if cls.has_scalar_variable(key): cls.del_scalar_variable(key) elif cls.has_array_variable(key): cls.del_array_variable(key)
def _core_wavefunction_has_variable(cls: core.Wavefunction, key: str) -> bool: """Whether scalar or array QCVariable *key* has been set on *self* :class:`psi4.core.Wavefunction`.""" return cls.has_scalar_variable(key) or cls.has_array_variable(key)