def _get_dos(self): """Return a doscar array node wrapped in a dictionary.""" vrp = self._parsers['vasprun.xml'] dcp = self._parsers['DOSCAR'] if not vrp or not dcp: return {'dos': None} dosnode = DataFactory('array')() # vrp.pdos is a numpy array, and thus not directly bool-convertible if vrp.pdos.size > 0: pdos = vrp.pdos.copy() for i, name in enumerate(vrp.pdos.dtype.names[1:]): num_spins = vrp.pdos.shape[1] # ~ pdos[name] = dcp[:, :, i+1:i+1+ns].transpose(0,2,1) cur = dcp.pdos[:, :, i + 1:i + 1 + num_spins].transpose(0, 2, 1) cond = vrp.pdos[name] < 0.1 pdos[name] = np.where(cond, cur, vrp.pdos[name]) dosnode.set_array('pdos', pdos) num_spins = 1 if dcp.tdos.shape[1] == 5: num_spins = 2 tdos = vrp.tdos[:num_spins, :].copy() for i, name in enumerate(vrp.tdos.dtype.names[1:]): cur = dcp.tdos[:, i + 1:i + 1 + num_spins].transpose() cond = vrp.tdos[:num_spins, :][name] < 0.1 tdos[name] = np.where(cond, cur, vrp.tdos[:num_spins, :][name]) dosnode.set_array('tdos', tdos) return {'dos': dosnode}
def get_dos_node(self, vrp, dcp): ''' takes VasprunParser and DosParser objects and returns a doscar array node ''' if not vrp or not dcp: return None dosnode = DataFactory('array')() if len(vrp.pdos): pdos = vrp.pdos.copy() for i, name in enumerate(vrp.pdos.dtype.names[1:]): ns = vrp.pdos.shape[1] # ~ pdos[name] = dcp[:, :, i+1:i+1+ns].transpose(0,2,1) cur = dcp.pdos[:, :, i + 1:i + 1 + ns].transpose(0, 2, 1) cond = vrp.pdos[name] < 0.1 pdos[name] = np.where(cond, cur, vrp.pdos[name]) dosnode.set_array('pdos', pdos) ns = 1 if dcp.tdos.shape[1] == 5: ns = 2 tdos = vrp.tdos[:ns, :].copy() for i, name in enumerate(vrp.tdos.dtype.names[1:]): cur = dcp.tdos[:, i + 1:i + 1 + ns].transpose() cond = vrp.tdos[:ns, :][name] < 0.1 tdos[name] = np.where(cond, cur, vrp.tdos[:ns, :][name]) dosnode.set_array('tdos', tdos) return dosnode