text:
from dlpix_wire.generators.base import BaseDataGenerator
import h5py
import logging
import numpy as np
import pdb

class Gen3D_pix(BaseDataGenerator):
  """
    Creates a generator for a list of files
  """
  logger = logging.getLogger("pixgen.gen.gen3d_pix")

  def __init__(self, datapaths, datasetname,
               labelsetname, batch_size=10, middle=False):
    self._files = [ i for i in datapaths]
    import random
    random.shuffle(self._files)

    self._middle = middle
    self._dataset = datasetname
    self._labelset = labelsetname
    self.batch_size = batch_size
    self.file_index=0
    self.current_index=0
#    self.truth = ["eminus", "eplus", "proton", "pizero", "piplus", "piminus", "muminus",  "muplus",  "kplus", "gamma"]
    self._plot=0
    self.truth = ["e-", "pi0", "gamma","p","mu"]

    self.labelvec = np.zeros(5)

    self.logger.info("Initializing npy file object with value: {}".format(self._files[self.current_index]))

    self.file_index = int(np.random.randint(len(self._files), size=1))
#    self.current_file = h5py.File(self._files[self.file_index], 'r')

    self.current_file = np.load(self._files[self.file_index])
    self.current_index = int(np.random.randint(self.current_file.shape[0], size=1))
    self.handlelabels(self.file_index)
    self.filterZeros()

  def filterZeros(self):
    # set low ADC pixels to 0
    pass

  def handlelabels(self,index):
    status = 0

#    x = self.current_file['image']['wires'] # from hdf5 days.
# Must bin the data evt-by-evt, since each event has different number of data elements (sim charge depositions)

    y = np.zeros((len(self.current_file['sdEvt']),250,600,250))
    ind = self.current_index

    self.logger.info("Binning data for event {}".format(ind))
    ### Should porbably do something here to mask out data that gets negative x's or lands in a non-central TPC. EC, 7-Jan-2018.

    data = np.array((self.current_file[ind,]['sdX'][self.current_file[ind,]['sdTPC']==3],self.current_file[ind,]['sdY'][self.current_file[ind,]['sdTPC']==3],self.current_file[ind,]['sdZ'][self.current_file[ind,]['sdTPC']==3] ))
    dataT = data.T

    if dataT.sum() is 0:
      print("Problem! Image is empty for ind " + str(ind))
      raise

    # dataT.shape => 3678, 3 e.g. meaning 3678 deositions.
    # Bail now if there aren't enough pixels with hits in them
    if len(self.current_file[ind,]['sdElec'][self.current_file[ind,]['sdTPC']==3]) < 500:
      self.logger.info("Bailing on this event with only {} energy depositions".format(len(self.current_file[ind,]['sdElec'][self.current_file[ind,]['sdTPC']==3])))
      status = 1
      return status

    xmin = np.argmin(dataT[:,0][dataT[:,0]>2])
    ymin = np.argmin(dataT[:,1][dataT[:,1]>2])
    zmin = np.argmin(dataT[:,2][dataT[:,2]>2])
    weights=self.current_file[ind,]['sdElec'][self.current_file[ind,]['sdTPC']==3]
                                    ##  view,chan,x
    H,edges = np.histogramdd(dataT,bins=(250,600,250),range=((0.,250.),(0.,600.),(0.,250.)),weights=weights)
#                    (Pdb) H.shape

    # Crop this back to central 2mx2mx2m about max activity point
    y = np.zeros((len(self.current_file['sdEvt']),200,200,200))
    (xmax,ymax,zmax) = np.unravel_index(np.argmax(H,axis=None),H.shape)
    # keep min past the minimal barrier, but not too close to far wall in each dimension.

    ix = np.maximum(xmax-100,0); ix = np.minimum(ix,250-200)
    iy = np.maximum(ymax-100,0); iy = np.minimum(iy,600-200)
    iz = np.maximum(zmax-100,0); iz = np.minimum(iz,250-200)

    y[ind,] = H[ix:ix+200,iy:iy+200,iz:iz+200]

    if self._plot<9:
      '''
      import matplotlib
      from mpl_toolkits.mplot3d import Axes3D
      matplotlib.use('Agg')
      import matplotlib.pyplot as plt

      fig = plt.figure()
      ax = fig.add_subplot(111, projection='3d')
      plt.imshow(y[ind,][2,])
      plt.colorbar()

      plt.savefig('scat_xy_'+str(ind)+'.png')
      ax.scatter(y[ind,])
      self._plot+=1
      plt.close()
      '''

      import matplotlib
      matplotlib.use('Agg')
      import matplotlib.pyplot as plt
      from mpl_toolkits.mplot3d import Axes3D
      fig = plt.figure(1)
      fig.clf()
      ax = Axes3D(fig)
      ax.scatter(dataT[:,0],dataT[:,1],dataT[:,0],c='b',marker='o', s=3) # s=np.log10(weights/np.mean(weights))+3 )
      ax.plot([], [], c='b', marker='o')
      ax.set_xlabel('X (drift)')
      ax.set_ylabel('Y (up))')
      ax.set_zlabel('Z (beam)')
      plt.draw()
      plt.savefig('scat_xyz_'+str(ind)+'.png')
      self._plot+=1
      plt.close()

    d = {}

    ptype = self._files[index].split("/")[-1].split("_")[-1].split(".")[0]
    # ptype = str(self._files[index].split("/")[-1].split("_")[0]) # "kplus", say
#    self._dataset = list(self.current_file.keys())[0]
#    self._labelset = list(self.current_file.keys())[1]

    # Setting these *set values were for hdf5 only to keep code from breaking ... I don't think these will be used elsewhere (EC, 7-Jan-2018)
    self._dataset = "image/pixels"
    self._labelset = "label/type"
    d[self._dataset] = y
    labelvectmp = np.array(self.labelvec)
#    print ("handle1evtfiles: particle and index are: " + str(ptype) + "  " + str(self.truth.index(ptype)))

    labelvectmp[self.truth.index(ptype)] = 1

    d[self._labelset] = np.expand_dims(labelvectmp,axis=0)
    self.current_file = d
    return status

  @property
  def output(self):
    current_index= self.current_index
    file_index = self.file_index

    x,y = next(self)
    self.current_index =current_index
    self.file_index = file_index
    return x[0].shape

  @property
  def input(self):
    current_index= self.current_index
    file_index = self.file_index
    x,y = next(self)
    self.current_index =current_index
    self.file_index = file_index
    return y[0].shape[0]

  def __len__(self):
    """
      Iterates over files to create the total sum length
      of the datasets in each file.
    """
    ### uggh. This is v slow, especially with tons of 1-event files.
#    return sum([h5py.File(i,'r')[self._dataset].shape[0] for i in self._files] )
    ### Let's just take a guess. Not sure we use this number anyway.
#    return len(self._files) * h5py.File(self._files[0],'r')[self._dataset].shape[0]
    return 12

  def __next__(self):
    return self.next()

  def debug_signal_handler(signal, frame):
    import pdb
    pdb.set_trace()
  import signal
#  signal.signal(signal.SIGINT, debug_signal_handler)

  def next(self):
    """
      This should iterate over both files and datasets within a file.
    """

    import pdb

    multifile = False

## This line and x=np.ndarray(1,1,tmp_x...) below need to be uncommented for 3D models, not so for 2D.

    xapp = np.empty(np.append(0,np.append(1,self.current_file[self._dataset].shape[1:])))
#    xapp = np.empty(np.append(1,self.current_file[self._dataset].shape[1:]))
    yapp = np.empty(np.append(0,self.current_file[self._labelset].shape[1:]))
    nevts = 0

    while self.batch_size > nevts:
      self.file_index = int(np.random.randint(len(self._files), size=1))
#      self.current_file = h5py.File(self._files[self.file_index], 'r')
      self.current_file = np.load(self._files[self.file_index])
      self.current_index = int(np.random.randint(self.current_file.shape[0], size=1))
      if self.handlelabels(self.file_index):  # This returns non-zero if any total image electron deposition is too low
        continue
      self.filterZeros()

      self.logger.info("Reading npy file: {}".format(self._files[self.file_index]))
      multifile = True

      # self.current_file not in fact a file at this point! handlelabels() turned it into an array.
      tmp_x =  self.current_file[self._dataset][self.current_index]

#      x = np.ndarray(shape=(1, tmp_x.shape[0],  tmp_x.shape[1],  tmp_x.shape[2]))
      x = np.ndarray(shape=(1, 1, tmp_x.shape[0],  tmp_x.shape[1],  tmp_x.shape[2]))
      x[0] = tmp_x
      y = self.current_file[self._labelset]

      if len(x) == 0 or len(y)==0 or not len(x) == len(y):
        return next(self)

      xapp = np.append(xapp,x,axis=0)
      yapp = np.append(yapp,y,axis=0)

      nevts += xapp.shape[1]

    if multifile:
      return (xapp,yapp)

    tmp_x = self.current_file[self._dataset][self.current_index:self.current_index+self.batch_size]
    x = np.ndarray(shape=(1, tmp_x.shape[0],  tmp_x.shape[1],  tmp_x.shape[2]))
    x[0] = tmp_x
    y = self.current_file[self._labelset]

    if len(x) == 0 or len(y)==0 or not len(x) == len(y):
      return next(self)
    return (x,y)

max_stars_repo_name: echurch/dlpix

max_issues_repo_name: echurch/dlpix

max_forks_repo_name: echurch/dlpix


text:
from typing import Any, Dict, List, Optional, Tuple, Type, Union

import gym
import numpy as np
from stable_baselines3.common.distributions import SquashedDiagGaussianDistribution
import torch as th
from torch.distributions.multivariate_normal import MultivariateNormal
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_

from scipy.optimize import minimize

from stable_baselines3.common import logger
from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise
from custom_algos.d3pg.off_policy_algorithm import OffPolicyAlgorithm
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
from stable_baselines3.common.utils import polyak_update
from custom_algos.d3pg.policies import MlpPolicy

class D3PG(OffPolicyAlgorithm):
    """
    Distributional Deterministic Deep Policy Gradient (D3PG)

    :param policy: The policy model to use (MlpPolicy, CnnPolicy, ...)
    :param env: The environment to learn from (if registered in Gym, can be str)
    :param learning_rate: learning rate for adam optimizer,
        the same learning rate will be used for all networks (Q-Values, Actor and Value function)
        it can be a function of the current progress remaining (from 1 to 0)
    :param buffer_size: size of the replay buffer
    :param learning_starts: how many steps of the model to collect transitions for before learning starts
    :param batch_size: Minibatch size for each gradient update
    :param tau: the soft update coefficient ("Polyak update", between 0 and 1)
    :param gamma: the discount factor
    :param train_freq: Update the model every ``train_freq`` steps. Alternatively pass a tuple of frequency and unit
        like ``(5, "step")`` or ``(2, "episode")``.
    :param gradient_steps: How many gradient steps to do after each rollout (see ``train_freq``)
        Set to ``-1`` means to do as many gradient steps as steps done in the environment
        during the rollout.
    :param action_noise: the action noise type (None by default), this can help
        for hard exploration problem. Cf common.noise for the different action noise type.
    :param optimize_memory_usage: Enable a memory efficient variant of the replay buffer
        at a cost of more complexity.
        See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195
    :param target_update_interval: update the target network every ``target_network_update_freq``
        gradient steps.
    :param create_eval_env: Whether to create a second environment that will be
        used for evaluating the agent periodically. (Only available when passing string for the environment)
    :param policy_kwargs: additional arguments to be passed to the policy on creation
    :param verbose: the verbosity level: 0 no output, 1 info, 2 debug
    :param seed: Seed for the pseudo random generators
    :param device: Device (cpu, cuda, ...) on which the code should be run.
        Setting it to auto, the code will be run on the GPU if possible.
    :param _init_setup_model: Whether or not to build the network at the creation of the instance
    """

    def __init__(
        self,
        policy: Union[str, Type[MlpPolicy]],
        env: Union[GymEnv, str],
        learning_rate: Union[float, Schedule] = 3e-4,
        buffer_size: int = int(1e6),
        learning_starts: int = 5000,
        batch_size: int = 256,
        tau: float = 0.02,
        gamma: float = 0.99,
        train_freq: Union[int, Tuple[int, str]] = 1,
        gradient_steps: int = 4,
        weight_decay: float = 0,
        action_dist_samples: int = 32,
        action_noise: Optional[ActionNoise] = None,
        optimize_memory_usage: bool = False,
        target_update_interval: int = 2,
        tensorboard_log: Optional[str] = None,
        create_eval_env: bool = False,
        policy_kwargs: Dict[str, Any] = None,
        verbose: int = 0,
        seed: Optional[int] = 0,
        device: Union[th.device, str] = "auto",
        _init_setup_model: bool = True,
    ):

        super(D3PG, self).__init__(
            policy,
            env,
            MlpPolicy,
            learning_rate,
            buffer_size,
            learning_starts,
            batch_size,
            tau,
            gamma,
            train_freq,
            gradient_steps,
            action_noise,
            policy_kwargs=policy_kwargs,
            tensorboard_log=tensorboard_log,
            verbose=verbose,
            device=device,
            create_eval_env=create_eval_env,
            seed=seed,
            optimize_memory_usage=optimize_memory_usage,
            supported_action_spaces=(gym.spaces.Box),
        )

        self.target_update_interval = target_update_interval
        self.weight_decay = weight_decay
        self.action_dist_samples = action_dist_samples

        if seed is not None:
            np.random.seed(seed)

        if _init_setup_model:
            self._setup_model()

    def _setup_model(self) -> None:
        super(D3PG, self)._setup_model()
        self._create_aliases()

    def _create_aliases(self) -> None:
        self.actor = self.policy.actor
        self.actor_target = self.policy.actor_target
        self.critic = self.policy.critic
        self.critic_target = self.policy.critic_target

        self.features_dim = self.actor.features_dim
        self.action_dim = self.actor.action_dim

    def train(self, gradient_steps: int, batch_size: int = 64) -> None:
        # Update optimizers learning rate
        optimizers = [self.actor.optimizer, self.critic.optimizer]

        # Update learning rate according to lr schedule
        self._update_learning_rate(optimizers)

        actor_losses, critic_losses = [], []

        for gradient_step in range(gradient_steps):
            # Sample replay buffer
            replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
            batch_size = replay_data.observations.size(0)

            # Critic update
            with th.no_grad():
                target_next_actions = self.actor_target.forward(replay_data.next_observations)
                target_next_actions_q, _ = self.critic_target.forward(
                    replay_data.next_observations, target_next_actions, self.action_dist_samples)
                target_next_actions_q = target_next_actions_q.transpose(1, 2)
                target_expected_Q = replay_data.rewards.unsqueeze(-1) + \
                    (1 - replay_data.dones.unsqueeze(-1)) * self.gamma * target_next_actions_q

            expected_Q, taus = self.critic.forward(replay_data.observations, replay_data.actions, self.action_dist_samples)

            # Quantile Huber loss
            td_error = target_expected_Q - expected_Q
            huber_1 = calculate_huber_loss(td_error, 1.0)
            quantil_1 = abs(taus - (td_error.detach() < 0).float()) * huber_1 / 1.0

            critic_loss = (quantil_1.sum(dim=1).mean(dim=1, keepdim=True) * replay_data.weights).mean()
            critic_losses.append(critic_loss.item())

            # Optimize critic
            self.critic.optimizer.zero_grad()
            critic_loss.backward()
            clip_grad_norm_(self.critic.parameters(), 1)
            self.critic.optimizer.step()

            # Actor update
            actions = self.actor.forward(replay_data.observations)
            actions_q = self.critic.get_qvalues(replay_data.observations, actions)
            actor_loss = -actions_q.mean()

            self.actor.optimizer.zero_grad()
            actor_loss.backward()
            self.actor.optimizer.step()

            actor_losses.append(actor_loss.item())

            self.replay_buffer.update_priorities(replay_data.indices, np.clip(
                abs(td_error.sum(dim=1).mean(dim=1, keepdim=True).data.cpu().numpy()), -1, 1))

            if gradient_step % self.target_update_interval == 0:
                polyak_update(self.actor.parameters(), self.actor_target.parameters(), self.tau)
                polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau)

        self._n_updates += gradient_steps

        logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
        logger.record("train/actor_loss", np.mean(actor_losses))
        logger.record("train/critic_loss", np.mean(critic_losses))

    def learn(
        self,
        total_timesteps: int,
        callback: MaybeCallback = None,
        log_interval: int = 4,
        eval_env: Optional[GymEnv] = None,
        eval_freq: int = -1,
        n_eval_episodes: int = 5,
        tb_log_name: str = "SAC",
        eval_log_path: Optional[str] = None,
        reset_num_timesteps: bool = True,
    ) -> OffPolicyAlgorithm:

        return super(D3PG, self).learn(
            total_timesteps=total_timesteps,
            callback=callback,
            log_interval=log_interval,
            eval_env=eval_env,
            eval_freq=eval_freq,
            n_eval_episodes=n_eval_episodes,
            tb_log_name=tb_log_name,
            eval_log_path=eval_log_path,
            reset_num_timesteps=reset_num_timesteps,
        )

    def _excluded_save_params(self) -> List[str]:
        return super(D3PG, self)._excluded_save_params() + ["actor", "actor_target", "critic", "critic_target"]

    def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
        state_dicts = ["policy", "actor.optimizer", "critic.optimizer"]
        saved_pytorch_variables = []
        # saved_pytorch_variables = ["log_ent_coef"]
        return state_dicts, saved_pytorch_variables

def calculate_huber_loss(td_errors, k=1.0):
    """
    Calculate huber loss element-wisely depending on kappa k.
    """
    loss = th.where(td_errors.abs() <= k, 0.5 * td_errors.pow(2), k * (td_errors.abs() - 0.5 * k))
    assert loss.shape == (td_errors.shape[0], 32, 32), "huber loss has wrong shape"
    return loss

# def get_min_qValues(critic_pairs):
#     critics = [pair[0] for pair in critic_pairs]
#     return th.min(th.cat(critics, dim=1), dim=1, keepdim=True).values

def get_huber_loss(target_q, q, taus, weights):
    td_error = target_q - q
    huber_1 = calculate_huber_loss(td_error, 1.0)
    quantil_1 = abs(taus - (td_error.detach() < 0).float()) * huber_1 / 1.0

    return (quantil_1.sum(dim=1).mean(dim=1, keepdim=True) * weights).mean()

ext: py

reason: import numpy,from scipy


text:
import os
import unittest
from shutil import rmtree

import h5py
import numpy as np
import torch

from torch_em import default_segmentation_loader
from torch_em.loss import DiceLoss
from torch_em.model import UNet2d

class TestDefaultTrainer(unittest.TestCase):
    data_path = './data.h5'
    checkpoint_folder = './checkpoints'
    log_folder = './logs'
    name = 'test'

    def setUp(self):
        shape = (8, 128, 128)
        chunks = (1, 128, 128)
        with h5py.File(self.data_path, 'w') as f:
            f.create_dataset('raw', data=np.random.rand(*shape), chunks=chunks)
            f.create_dataset('labels', data=np.random.randint(0, 32, size=shape), chunks=chunks)

    def tearDown(self):
        if os.path.exists(self.data_path):
            os.remove(self.data_path)
        if os.path.exists(self.checkpoint_folder):
            rmtree(self.checkpoint_folder)
        if os.path.exists(self.log_folder):
            rmtree(self.log_folder)

    def _get_kwargs(self):
        loader = default_segmentation_loader(
            raw_paths=self.data_path, raw_key='raw',
            label_paths=self.data_path, label_key='labels',
            batch_size=1, patch_shape=(1, 128, 128), ndim=2
        )
        model = UNet2d(in_channels=1, out_channels=1,
                       depth=2, initial_features=4)
        kwargs = {
            'name': self.name,
            'train_loader': loader,
            'val_loader': loader,
            'model': model,
            'loss': DiceLoss(),
            'metric': DiceLoss(),
            'optimizer': torch.optim.Adam(model.parameters(), lr=1e-5),
            'device': torch.device('cpu'),
            'mixed_precision': False
        }
        return kwargs

    def test_fit(self):
        from torch_em.trainer import DefaultTrainer
        trainer = DefaultTrainer(**self._get_kwargs())
        trainer.fit(10)

        save_folder = os.path.join(self.checkpoint_folder, self.name)
        self.assertTrue(os.path.exists(save_folder))
        best_ckpt = os.path.join(save_folder, 'best.pt')
        self.assertTrue(os.path.exists(best_ckpt))
        latest_ckpt = os.path.join(save_folder, 'latest.pt')
        self.assertTrue(os.path.exists(latest_ckpt))
        self.assertEqual(trainer.iteration, 10)

        trainer.fit(2)
        self.assertEqual(trainer.iteration, 12)

        trainer = DefaultTrainer(**self._get_kwargs())
        trainer.fit(8, load_from_checkpoint='latest')
        self.assertEqual(trainer.iteration, 20)

    def test_from_checkpoint(self):
        from torch_em.trainer import DefaultTrainer
        trainer = DefaultTrainer(**self._get_kwargs())
        trainer.fit(10)

        trainer2 = DefaultTrainer.from_checkpoint(
            os.path.join(self.checkpoint_folder, self.name),
            name='latest'
        )
        self.assertEqual(trainer.iteration, trainer2.iteration)

        # make sure that the optimizer was loaded properly
        lr1 = [pm['lr'] for pm in trainer.optimizer.param_groups][0]
        lr2 = [pm['lr'] for pm in trainer2.optimizer.param_groups][0]
        self.assertEqual(lr1, lr2)

        trainer2.fit(10)
        self.assertEqual(trainer2.iteration, 20)

if __name__ == '__main__':
    unittest.main()

max_stars_repo_name: oeway/torch-em

max_issues_repo_name: oeway/torch-em

max_forks_repo_name: oeway/torch-em

reason: import numpy


text:
# coding: utf-8
"""
Parses the computed data from VASP defect calculations.
"""
# from __future__ import unicode_literals
from __future__ import division

import glob
import logging
import os
import warnings

import numpy as np
from monty.json import MontyDecoder
from monty.serialization import loadfn
from pymatgen.analysis.defects.core import Vacancy, Substitution, Interstitial, DefectEntry
from pymatgen.analysis.defects.defect_compatibility import DefectCompatibility
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.core import PeriodicSite, Structure
from pymatgen.entries.computed_entries import ComputedStructureEntry
from pymatgen.ext.matproj import MPRester
from pymatgen.io.vasp.inputs import Potcar, UnknownPotcarWarning
from pymatgen.io.vasp.outputs import Vasprun, Locpot, Outcar, Poscar

from doped.pycdt.core import chemical_potentials

def convert_cd_to_de(cd, b_cse):
    """
    As of pymatgen v2.0, ComputedDefect objects were deprecated in favor
    of DefectEntry objects in pymatgen.analysis.defects.core
    This function takes a ComputedDefect (either as a dict or object) and
    converts it into a DefectEntry object in order to handle legacy
    PyCDT creation within the current paradigm of PyCDT.

    :param cd (dict or ComputedDefect object): ComputedDefect as an object or as a dictionary
    :params b_cse (dict or ComputedStructureEntry object): ComputedStructureEntry of bulk entry
        associated with the ComputedDefect.
    :return: de (DefectEntry): Resulting DefectEntry object
    """
    if type(cd) != dict:
        cd = cd.as_dict()
    if type(b_cse) != dict:
        b_cse = b_cse.as_dict()

    bulk_sc_structure = Structure.from_dict(b_cse["structure"])

    # modify defect_site as required for Defect object, confirming site exists in bulk structure
    site_cls = cd["site"]
    defect_site = PeriodicSite.from_dict(site_cls)
    def_nom = cd["name"].lower()
    if "sub_" in def_nom or "as_" in def_nom:
        # modify site object for substitution site of Defect object
        site_cls["species"][0]["element"] = cd["name"].split("_")[2]
        defect_site = PeriodicSite.from_dict(site_cls)

    poss_deflist = sorted(
        bulk_sc_structure.get_sites_in_sphere(defect_site.coords, 0.2, include_index=True),
        key=lambda x: x[1],
    )
    if len(poss_deflist) != 1:
        raise ValueError(
            "ComputedDefect to DefectEntry conversion failed. "
            "Could not determine periodic site position in bulk supercell."
        )

    # create defect object
    if "vac_" in def_nom:
        defect_obj = Vacancy(bulk_sc_structure, defect_site, charge=cd["charge"])
    elif "as_" in def_nom or "sub_" in def_nom:
        defect_obj = Substitution(bulk_sc_structure, defect_site, charge=cd["charge"])
    elif "int_" in def_nom:
        defect_obj = Interstitial(bulk_sc_structure, defect_site, charge=cd["charge"])
    else:
        raise ValueError("Could not recognize defect type for {}".format(cd["name"]))

    # assign proper energy and parameter metadata
    uncorrected_energy = cd["entry"]["energy"] - b_cse["energy"]
    def_path = os.path.split(cd["entry"]["data"]["locpot_path"])[0]
    bulk_path = os.path.split(b_cse["data"]["locpot_path"])[0]
    p = {"defect_path": def_path, "bulk_path": bulk_path, "encut": cd["entry"]["data"]["encut"]}

    de = DefectEntry(defect_obj, uncorrected_energy, parameters=p)

    return de

def get_vasprun(vasprun_path, **kwargs):
    """ Read the vasprun.xml(.gz) file as a pymatgen Locpot object """
    warnings.filterwarnings(
        "ignore", category=UnknownPotcarWarning
    )  # Ignore POTCAR warnings when loading vasprun.xml
    # pymatgen assumes the default PBE with no way of changing this within get_vasprun())
    warnings.filterwarnings("ignore", message="No POTCAR file with matching TITEL fields")
    if os.path.exists(vasprun_path):
        vasprun = Vasprun(vasprun_path)
    elif os.path.exists(vasprun_path + ".gz", **kwargs):
        vasprun = Vasprun(vasprun_path + ".gz", **kwargs)
    else:
        raise FileNotFoundError(
            f"""I can't find a vasprun.xml(.gz) at {vasprun_path}(.gz).
                   You sure there's one there pal? I need it to parse the calculation results"""
        )
    return vasprun

def get_locpot(locpot_path):
    """ Read the LOCPOT(.gz) file as a pymatgen Locpot object """
    if os.path.exists(locpot_path):
        locpot = Locpot.from_file(locpot_path)
    elif os.path.exists(locpot_path + ".gz"):
        locpot = Locpot.from_file(locpot_path + ".gz")
    else:
        raise FileNotFoundError(
            f"""I can't find a LOCPOT(.gz) at {locpot_path}(.gz).
                   You sure there's one there pal? I need it to get the Freysoldt correction"""
        )
    return locpot

class SingleDefectParser:
    def __init__(
        self,
        defect_entry,
        compatibility=DefectCompatibility(plnr_avg_var_tol=0.0002),
        defect_vr=None,
        bulk_vr=None,
    ):
        """
        Parse a defect object using features that resemble that of a standard
        DefectBuilder object (emmet), but without the requirement of atomate.
        Also allows for use of DefectCompatibility object within pymatgen

        :param defect_entry (DefectEntry): DefectEntry of interest (using the bulk supercell as
        bulk_structure)
            NOTE: to make use of methods within the class, bulk_path and and defect_path
            must exist within the defect_entry parameters class.
        :param compatibility (DefectCompatibility): Compatibility class instance for
            performing compatibility analysis on defect entry.
        :param defect_vr (Vasprun):
        :param bulk_vr (Vasprun):

        """
        self.defect_entry = defect_entry
        self.compatibility = compatibility
        self.defect_vr = defect_vr
        self.bulk_vr = bulk_vr

    @staticmethod
    def from_paths(
        path_to_defect,
        path_to_bulk,
        dielectric,
        defect_charge,
        mpid=None,
        compatibility=DefectCompatibility(plnr_avg_var_tol=0.0002),
        initial_defect_structure=None,
    ):
        """
        Identify defect object based on file paths. Minimal parsing performing for
        instantiating the SingleDefectParser class.

        :param path_to_defect (str): path to defect folder of interest (with vasprun.xml(.gz))
        :param path_to_bulk (str): path to bulk folder of interest (with vasprun.xml(.gz))
        :param dielectric (float or 3x3 matrix): ionic + static contributions to dielectric constant
        :param defect_charge (int):
        :param mpid (str):
        :param compatibility (DefectCompatibility): Compatibility class instance for
            performing compatibility analysis on defect entry.

        Return:
            Instance of the SingleDefectParser class.
        """
        parameters = {
            "bulk_path": path_to_bulk,
            "defect_path": path_to_defect,
            "dielectric": dielectric,
            "mpid": mpid,
        }

        # add bulk simple properties
        bulk_vr = get_vasprun(os.path.join(path_to_bulk, "vasprun.xml"))
        bulk_energy = bulk_vr.final_energy
        bulk_sc_structure = bulk_vr.initial_structure.copy()

        # add defect simple properties
        defect_vr = get_vasprun(os.path.join(path_to_defect, "vasprun.xml"))
        defect_energy = defect_vr.final_energy
        # Can specify initial defect structure (to help PyCDT find the defect site if
        # multiple relaxations were required, else use from defect relaxation OUTCAR:
        if initial_defect_structure:
            initial_defect_structure = Poscar.from_file(initial_defect_structure).structure.copy()
        else:
            initial_defect_structure = defect_vr.initial_structure.copy()

        # Add initial defect structure to parameters, so it can be pulled later on (eg. for Kumagai loader)
        parameters["initial_defect_structure"] = initial_defect_structure

        # identify defect site, structural information, and create defect object
        num_ids = len(initial_defect_structure)
        num_bulk = len(bulk_sc_structure)
        if num_ids == num_bulk - 1:
            defect_type = "Vacancy"
        elif num_ids == num_bulk + 1:
            defect_type = "Interstitial"
        elif num_ids == num_bulk:
            defect_type = "Substitution"
        else:
            raise ValueError(
                "Could not identify defect type just from number of sites in structure: "
                "{} in bulk vs. {} in defect?".format(num_bulk, num_ids)
            )

        defect_index_sc_coords = None
        transformation_path = os.path.join(path_to_defect, "transformation.json")
        if not os.path.exists(transformation_path): # try next folder up
            orig_transformation_path = transformation_path
            transformation_path = os.path.join(
                os.path.dirname(os.path.normpath(path_to_defect)), "transformation.json"
            )
            if os.path.exists(transformation_path):
                print("No transformation file found at {}, but found one at {}. "
                      "Using this for defect parsing.".format(orig_transformation_path, transformation_path))

        if os.path.exists(transformation_path):
            tf = loadfn(transformation_path)
            site = tf["defect_supercell_site"]
            if defect_type == "Vacancy": # should we also search the bulk structure for substitutions?
                poss_deflist = sorted(
                    bulk_sc_structure.get_sites_in_sphere(site.coords, 0.2, include_index=True),
                    key=lambda x: x[1],
                )
                searched = "bulk_supercell"
            else:
                poss_deflist = sorted(
                    initial_defect_structure.get_sites_in_sphere(
                        site.coords, 0.2, include_index=True
                    ),
                    key=lambda x: x[1],
                )
                searched = "initial_defect_structure"
            if not poss_deflist:
                raise ValueError(
                    "{} specified defect site {}, but could not find it in {}."
                    " Abandoning parsing".format(transformation_path, site, searched)
                )
            defect_index_sc_coords = poss_deflist[0][2]
        else:
            print(
                "No transformation file exists at {} or {}.\nCalculating defect index manually"
                " (proceed with caution)".format(orig_transformation_path, transformation_path)
            )

        # IF not transformation file exists, the defect_index_sc_coords will not be identified in
        # previous routine, proceed by identifying the defect site through a comparison of bulk
        # sites and initial defect structure sites.
        # WARNING: this can cause issues if initial_defect_structure is slightly different than
        # bulk_sc_structure (as a result of multiple relaxation steps, for example)
        # noinspection DuplicatedCode
        if defect_index_sc_coords is None:
            bulksites = [site.frac_coords for site in bulk_sc_structure]
            initsites = [site.frac_coords for site in initial_defect_structure]
            distmatrix = initial_defect_structure.lattice.get_all_distances(bulksites, initsites)
            min_dist_with_index = [
                [min(distmatrix[bulk_index]), int(bulk_index), int(distmatrix[bulk_index].argmin())]
                for bulk_index in range(len(distmatrix))
            ]  # list of [min dist, bulk ind, defect ind]

            site_matching_indices = []
            poss_defect = []
            if defect_type in ["Vacancy", "Interstitial"]:
                for mindist, bulk_index, defect_index in min_dist_with_index:
                    if mindist < 0.2:
                        site_matching_indices.append([bulk_index, defect_index])
                    elif defect_type == "Vacancy":
                        poss_defect.append([bulk_index, bulksites[bulk_index][:]])

                if defect_type == "Interstitial":
                    poss_defect = [
                        [ind, fc[:]]
                        for ind, fc in enumerate(initsites)
                        if ind not in np.array(site_matching_indices)[:, 1]
                    ]

            elif defect_type == "Substitution":
                # noinspection DuplicatedCode
                for mindist, bulk_index, defect_index in min_dist_with_index:
                    species_match = (
                        bulk_sc_structure[bulk_index].specie
                        == initial_defect_structure[defect_index].specie
                    )
                    if mindist < 0.2 and species_match:
                        site_matching_indices.append([bulk_index, defect_index])

                    elif not species_match:
                        poss_defect.append([defect_index, initsites[defect_index][:]])

            if len(poss_defect) == 1:
                defect_index_sc_coords = poss_defect[0][0]
            else:
                raise ValueError(
                    "Found {} possible defect sites when matching bulk and "
                    "defect structure".format(len(poss_defect))
                )

            if len(set(np.array(site_matching_indices)[:, 0])) != len(
                set(np.array(site_matching_indices)[:, 1])
            ):
                raise ValueError(
                    "Error occurred in site_matching routine. Double counting of site matching "
                    "occurred:{}\nAbandoning structure parsing.".format(site_matching_indices)
                )

        if defect_type == "Vacancy":
            defect_site = bulk_sc_structure[defect_index_sc_coords]
        else:
            defect_site = initial_defect_structure[defect_index_sc_coords]

        for_monty_defect = {
            "@module": "pymatgen.analysis.defects.core",
            "@class": defect_type,
            "charge": defect_charge,
            "structure": bulk_sc_structure,
            "defect_site": defect_site,
        }
        defect = MontyDecoder().process_decoded(for_monty_defect)
        test_defect_structure = defect.generate_defect_structure()
        if not StructureMatcher(
            stol=0.5, primitive_cell=False, scale=False, attempt_supercell=False, allow_subset=False
        ).fit(test_defect_structure, defect_vr.initial_structure):
            # NOTE: this does not insure that cartesian coordinates or indexing are identical
            # Note: I've changed stol to 0.5 to fix matching for defects that distort significantly
            if not StructureMatcher(
                stol=0.5,
                primitive_cell=False,
                scale=False,
                attempt_supercell=False,
                allow_subset=False,
            ).fit(test_defect_structure, initial_defect_structure):
                # defect_vr.initial_structure may not match test_defect_structure if we've had
                # significant relaxation prior to final defect energy calculation, so test against
                # initial_defect_structure (which may have been specified in function arguments)
                raise ValueError("Error in defect object matching!")

        defect_entry = DefectEntry(
            defect, defect_energy - bulk_energy, corrections={}, parameters=parameters
        )

        return SingleDefectParser(
            defect_entry, compatibility=compatibility, defect_vr=defect_vr, bulk_vr=bulk_vr
        )

    def freysoldt_loader(self, bulk_locpot=None):
        """Load metadata required for performing Freysoldt correction
        requires "bulk_path" and "defect_path" to be loaded to DefectEntry parameters dict.
        Can read gunzipped "LOCPOT.gz" files as well.

        Args:
            bulk_locpot (Locpot): Add bulk Locpot object for expedited parsing.
                If None, will load from file path variable bulk_path
        Return:
            bulk_locpot object for reuse by another defect entry (for expedited parsing)
        """
        if not self.defect_entry.charge:
            # dont need to load locpots if charge is zero
            return None

        if not bulk_locpot:
            bulk_locpot_path = os.path.join(self.defect_entry.parameters["bulk_path"], "LOCPOT")
            bulk_locpot = get_locpot(bulk_locpot_path)

        def_locpot_path = os.path.join(self.defect_entry.parameters["defect_path"], "LOCPOT")
        def_locpot = get_locpot(def_locpot_path)

        axis_grid = [def_locpot.get_axis_grid(i) for i in range(3)]
        bulk_planar_averages = [bulk_locpot.get_average_along_axis(i) for i in range(3)]
        defect_planar_averages = [def_locpot.get_average_along_axis(i) for i in range(3)]

        defect_frac_sc_coords = self.defect_entry.site.frac_coords

        self.defect_entry.parameters.update(
            {
                "axis_grid": axis_grid,
                "bulk_planar_averages": bulk_planar_averages,
                "defect_planar_averages": defect_planar_averages,
                "initial_defect_structure": def_locpot.structure,
                "defect_frac_sc_coords": defect_frac_sc_coords,
            }
        )

        return bulk_locpot

    # noinspection DuplicatedCode
    def kumagai_loader(self, bulk_outcar=None):
        """Load metadata required for performing Kumagai correction
        requires "bulk_path" and "defect_path" to be loaded to DefectEntry parameters dict.

        Args:
            bulk_outcar (Outcar): Add bulk Outcar object for expedited parsing.
                If None, will load from file path variable bulk_path
        Return:
            bulk_outcar object for reuse by another defect entry (for expedited parsing)
        """
        if not self.defect_entry.charge:
            # dont need to load outcars if charge is zero
            return None

        if not bulk_outcar:
            bulk_outcar_path = os.path.join(self.defect_entry.parameters["bulk_path"], "OUTCAR")
            bulk_outcar = Outcar(bulk_outcar_path)

        def_outcar_path = os.path.join(self.defect_entry.parameters["defect_path"], "OUTCAR")
        def_outcar = Outcar(def_outcar_path)

        bulk_atomic_site_averages = bulk_outcar.electrostatic_potential
        defect_atomic_site_averages = def_outcar.electrostatic_potential

        # Is parsing the bulk structure and defect site here again unnecessary?
        # Yes, if the user has used SingleDefectParser.from_paths, but not if they've
        # been a cowboy and used SingleDefectParser() with a pre-generated DefectEntry...
        # So let's keep it for the moment

        if os.path.exists(os.path.join(self.defect_entry.parameters["bulk_path"], "POSCAR")):
            bulk_sc_structure = Poscar.from_file(
                os.path.join(self.defect_entry.parameters["bulk_path"], "POSCAR")
            ).structure
        elif self.bulk_vr:
            bulk_sc_structure = self.bulk_vr.initial_structure.copy()
        else:
            bulk_sc_structure = get_vasprun(
                os.path.join(self.defect_entry.parameters["bulk_path"], "vasprun.xml")
            ).initial_structure.copy()

        if "initial_defect_structure" in self.defect_entry.parameters:
            initial_defect_structure = self.defect_entry.parameters["initial_defect_structure"]
        elif os.path.exists(os.path.join(self.defect_entry.parameters["defect_path"], "POSCAR")):
            initial_defect_structure = Poscar.from_file(
                os.path.join(self.defect_entry.parameters["defect_path"], "POSCAR")
            ).structure
        elif self.defect_vr:
            initial_defect_structure = self.defect_vr.initial_structure
        else:
            initial_defect_structure = get_vasprun(
                os.path.join(self.defect_entry.parameters["defect_path"], "vasprun.xml")
            ).initial_structure

        bulksites = [site.frac_coords for site in bulk_sc_structure]
        initsites = [site.frac_coords for site in initial_defect_structure]
        distmatrix = initial_defect_structure.lattice.get_all_distances(
            bulksites, initsites
        )  # first index of this list is bulk index
        min_dist_with_index = [
            [min(distmatrix[bulk_index]), int(bulk_index), int(distmatrix[bulk_index].argmin())]
            for bulk_index in range(len(distmatrix))
        ]  # list of [min dist, bulk ind, defect ind]

        site_matching_indices = []
        poss_defect = []
        defect_index_sc_coords = None
        if isinstance(self.defect_entry.defect, (Vacancy, Interstitial)):
            for mindist, bulk_index, defect_index in min_dist_with_index:
                if mindist < 0.2:
                    site_matching_indices.append([bulk_index, defect_index])
                elif (
                    isinstance(self.defect_entry.defect, Vacancy) and defect_index_sc_coords is None
                ):
                    poss_defect.append([bulk_index, bulksites[bulk_index][:]])

            if (
                isinstance(self.defect_entry.defect, Interstitial)
                and defect_index_sc_coords is None
            ):
                poss_defect = [
                    [ind, fc[:]]
                    for ind, fc in enumerate(initsites)
                    if ind not in np.array(site_matching_indices)[:, 1]
                ]

        elif isinstance(self.defect_entry.defect, Substitution):
            # noinspection DuplicatedCode
            for mindist, bulk_index, defect_index in min_dist_with_index:
                species_match = (
                    bulk_sc_structure[bulk_index].specie
                    == initial_defect_structure[defect_index].specie
                )
                if mindist < 0.2 and species_match:
                    site_matching_indices.append([bulk_index, defect_index])

                elif not species_match and defect_index_sc_coords is None:
                    poss_defect.append([defect_index, initsites[defect_index][:]])

        if defect_index_sc_coords is None:
            if len(poss_defect) == 1:
                defect_index_sc_coords = poss_defect[0][0]
                defect_frac_sc_coords = poss_defect[0][1]

            elif self.defect_entry.site:
                defect_frac_sc_coords = self.defect_entry.site.frac_coords
                defect_type = self.defect_entry.defect.as_dict()["@class"]
                if defect_type == "Vacancy":
                    poss_deflist = sorted(
                        bulk_sc_structure.get_sites_in_sphere(
                            self.defect_entry.site.coords, 0.2, include_index=True
                        ),
                        key=lambda x: x[1],
                    )
                    searched = "bulk_supercell"
                else:
                    poss_deflist = sorted(
                        initial_defect_structure.get_sites_in_sphere(
                            self.defect_entry.site.coords, 0.2, include_index=True
                        ),
                        key=lambda x: x[1],
                    )
                    searched = "initial_defect_structure"
                if not poss_deflist:
                    raise ValueError(
                        "Could not find defect site {} in {}, abandoning parsing".format(
                            self.defect_entry.site, searched)
                    )
                defect_index_sc_coords = poss_deflist[0][2]

            else:
                raise ValueError(
                    "Found {} possible defect sites when matching bulk and "
                    "defect structure".format(len(poss_defect))
                )

        if len(set(np.array(site_matching_indices)[:, 0])) != len(
            set(np.array(site_matching_indices)[:, 1])
        ):
            raise ValueError(
                "Error occurred in site_matching routine. Double counting of site matching "
                "occurred:{}\nAdvising against Kumagai parsing.".format(site_matching_indices)
            )

        # user Wigner-Seitz radius for sampling radius
        wz = initial_defect_structure.lattice.get_wigner_seitz_cell()
        dist = []
        for facet in wz:
            midpt = np.mean(np.array(facet), axis=0)
            dist.append(np.linalg.norm(midpt))
        sampling_radius = min(dist)

        self.defect_entry.parameters.update(
            {
                "bulk_atomic_site_averages": bulk_atomic_site_averages,
                "defect_atomic_site_averages": defect_atomic_site_averages,
                "initial_defect_structure": initial_defect_structure,
                "site_matching_indices": site_matching_indices,
                "sampling_radius": sampling_radius,
                "defect_frac_sc_coords": defect_frac_sc_coords,
                "defect_index_sc_coords": defect_index_sc_coords,
            }
        )

        return bulk_outcar

    def get_stdrd_metadata(self):

        if not self.bulk_vr:
            path_to_bulk = self.defect_entry.parameters["bulk_path"]
            self.bulk_vr = get_vasprun(os.path.join(path_to_bulk, "vasprun.xml"))

        if not self.defect_vr:
            path_to_defect = self.defect_entry.parameters["defect_path"]
            self.defect_vr = get_vasprun(os.path.join(path_to_defect, "vasprun.xml"))

        # standard bulk metadata
        bulk_energy = self.bulk_vr.final_energy
        bulk_sc_structure = self.bulk_vr.initial_structure
        self.defect_entry.parameters.update(
            {"bulk_energy": bulk_energy, "bulk_sc_structure": bulk_sc_structure}
        )

        # standard run metadata
        run_metadata = {}
        run_metadata.update(
            {
                "defect_incar": self.defect_vr.incar,
                "bulk_incar": self.bulk_vr.incar,
                "defect_kpoints": self.defect_vr.kpoints,
                "bulk_kpoints": self.bulk_vr.kpoints,
            }
        )
        run_metadata.update(
            {
                "incar_calctype_summary": {
                    k: self.defect_vr.incar.get(k, None)
                    if self.defect_vr.incar.get(k) not in ["None", "False", False]
                    else None
                    for k in ["LHFCALC", "HFSCREEN", "IVDW", "LUSE_VDW", "LDAU", "METAGGA"]
                }
            }
        )
        run_metadata.update(
            {
                "potcar_summary": {
                    "pot_spec": [potelt["titel"] for potelt in self.defect_vr.potcar_spec],
                    "pot_labels": self.defect_vr.potcar_spec,
                    "pot_type": self.defect_vr.run_type,
                }
            }
        )

        self.defect_entry.parameters.update({"run_metadata": run_metadata.copy()})

        # standard defect run metadata
        self.defect_entry.parameters.update(
            {
                "final_defect_structure": self.defect_vr.final_structure,
                "initial_defect_structure": self.defect_vr.initial_structure,
                "defect_energy": self.defect_vr.final_energy,
            }
        )

        # grab defect energy and eigenvalue information for band filling and localization analysis
        eigenvalues = {
            spincls.value: eigdict.copy() for spincls, eigdict in self.defect_vr.eigenvalues.items()
        }
        kpoint_weights = self.defect_vr.actual_kpoints_weights[:]
        self.defect_entry.parameters.update(
            {"eigenvalues": eigenvalues, "kpoint_weights": kpoint_weights}
        )

    def get_bulk_gap_data(self, no_MP=False, actual_bulk_path=None):

        if not self.bulk_vr:
            path_to_bulk = self.defect_entry.parameters["bulk_path"]
            self.bulk_vr = get_vasprun(os.path.join(path_to_bulk, "vasprun.xml"))

        bulk_sc_structure = self.bulk_vr.initial_structure
        mpid = self.defect_entry.parameters["mpid"]

        if not mpid and not no_MP:
            try:
                with MPRester() as mp:
                    tmp_mplist = mp.get_entries_in_chemsys(list(bulk_sc_structure.symbol_set))
                mplist = [
                    ment.entry_id
                    for ment in tmp_mplist
                    if ment.composition.reduced_composition
                    == bulk_sc_structure.composition.reduced_composition
                ]
            except:
                raise ValueError(
                    "Error with querying MPRester for {}"
                    "".format(bulk_sc_structure.composition.reduced_formula)
                )

            mpid_fit_list = []
            for trial_mpid in mplist:
                with MPRester() as mp:
                    mpstruct = mp.get_structure_by_material_id(trial_mpid)
                if StructureMatcher(
                    primitive_cell=True, scale=False, attempt_supercell=True, allow_subset=False
                ).fit(bulk_sc_structure, mpstruct):
                    mpid_fit_list.append(trial_mpid)

            if len(mpid_fit_list) == 1:
                mpid = mpid_fit_list[0]
                print("Single mp-id found for bulk structure:{}.".format(mpid))
            elif len(mpid_fit_list) > 1:
                num_mpid_list = [int(mp.split("" - "")[1]) for mp in mpid_fit_list]
                num_mpid_list.sort()
                mpid = "mp-" + str(num_mpid_list[0])
                print(
                    "Multiple mp-ids found for bulk structure:{}\nWill use lowest number mpid "
                    "for bulk band structure = {}.".format(str(mpid_fit_list), mpid)
                )
            else:
                print(
                    "Could not find bulk structure in MP database after tying the "
                    "following list:\n{}".format(mplist)
                )
                mpid = None
        else:
            print("Manually fed mpid = {}".format(mpid))

        vbm, cbm, bandgap = None, None, None
        gap_parameters = {}
        if mpid is not None and not no_MP:
            with MPRester() as mp:
                bs = mp.get_bandstructure_by_material_id(mpid)
            if bs:
                cbm = bs.get_cbm()["energy"]
                vbm = bs.get_vbm()["energy"]
                bandgap = bs.get_band_gap()["energy"]
                gap_parameters.update(
                    {"MP_gga_BScalc_data": bs.get_band_gap().copy()}
                )  # contains gap kpt transition

        if vbm is None or bandgap is None or cbm is None or no_MP or not actual_bulk_path:
            if mpid and bandgap is None:
                print(
                    "WARNING: Mpid {} was provided, but no bandstructure entry currently exists "
                    "for it. \n"
                    "Reverting to use of bulk supercell calculation for band edge extrema.".format(
                        mpid
                    )
                )
            if mpid and no_MP:
                print(
                    "Mpid {} was provided, but we've decided to say:\n"
                    "'Thanks, but no thanks' to the Materials Project.\n"
                    "Reverting to use of bulk supercell calculation for band edge extrema.".format(
                        mpid
                    )
                )
            else:
                print(
                    "WARNING: No mp-id provided, will fetch CBM/VBM details from the "
                    "bulk calculation."
                )
                print(
                    "This may not be appropriate if the VBM/CBM occur at reciprocal points "
                    "not included in the bulk calculation."
                )

            gap_parameters.update({"MP_gga_BScalc_data": None})  # to signal no MP BS is used
            bandgap, cbm, vbm, _ = self.bulk_vr.eigenvalue_band_properties

        # Note that we've modified
        # pycdt.utils.parse_calculations.SingleDefectParser.get_bulk_gap_data() to have
        # the option to say 'no_MP', where it will use the bulk calculation bandgap, regardless
        # of whether a shit GGA bandgap
        # for our material exists in the MP database.
        # We've also added actual_bulk optional argument to
        # pycdt.utils.parse_calculations.SingleDefectParser.get_bulk_gap_data(), in case we're
        # looking at defect complexes (so the SingleDefectParser 'bulk' is actually the original
        # point defect, with the wrong pmg-computed gap)

        if actual_bulk_path:
            print(f"Using actual bulk path: {actual_bulk_path}")
            actual_bulk_vr = get_vasprun(os.path.join(actual_bulk_path, "vasprun.xml"))
            bandgap, cbm, vbm, _ = actual_bulk_vr.eigenvalue_band_properties

        gap_parameters.update({"mpid": mpid, "cbm": cbm, "vbm": vbm, "gap": bandgap})
        self.defect_entry.parameters.update(gap_parameters)

    def run_compatibility(self):
        # Set potalign so pymatgen can calculate bandfilling for 'neutral' defects (possible for resonant dopants etc.)
        if self.defect_entry.charge == 0 and "potalign" not in self.defect_entry.parameters:
            self.defect_entry.parameters["potalign"] = 0
        self.defect_entry = self.compatibility.process_entry(self.defect_entry)
        if not self.defect_entry.parameters["is_compatible"]:
            delocalized_warning = f"""
Delocalization analysis has indicated that {self.defect_entry.name}
with charge {self.defect_entry.charge} may not be compatible with the chosen charge correction
scheme, and may require a larger supercell for accurate calculation of the energy. Recommended to
look at the correction plots (i.e. run `get_correction_freysoldt(DefectEntry,...,plot=True)` from
`doped.pycdt.corrections.finite_size_charge_correction`) to visually determine if
charge correction scheme still appropriate, then `sdp.compatibility.perform_freysoldt(DefectEntry)`
if you're happy (replace 'freysoldt' with 'kumagai' if using anisotropic correction).
You can also change the DefectCompatibility() tolerance settings via the `compatibility` parameter
in `SingleDefectParser.from_paths()`.
Watch out that if `num_hole_vbm` or `num_elec_cbm` are greater than the free_chg_cutoff (default
2.1), charge correction will not be applied."""
            warnings.warn(message=delocalized_warning, stacklevel=1)

class PostProcess:
    def __init__(self, root_fldr, mpid=None, mapi_key=None):
        """
        Post processing object for charged point-defect calculations.

        Args:
            root_fldr (str): path (relative) to directory
                in which data of charged point-defect calculations for
                a particular system are to be found;
            mpid (str): Materials Project ID of bulk structure;
                format "mp-X", where X is an integer;
            mapi_key (str): Materials API key to access database.

        """
        self._root_fldr = root_fldr
        self._mpid = mpid
        self._mapi_key = mapi_key
        self._substitution_species = set()

    def parse_defect_calculations(self):
        """
        Parses the defect calculations as DefectEntry objects,
        from a PyCDT root_fldr file structure.
        Charge correction is missing in the first run.
        """
        logger = logging.getLogger(__name__)
        parsed_defects = []
        subfolders = glob.glob(os.path.join(self._root_fldr, "vac_*"))
        subfolders += glob.glob(os.path.join(self._root_fldr, "as_*"))
        subfolders += glob.glob(os.path.join(self._root_fldr, "sub_*"))
        subfolders += glob.glob(os.path.join(self._root_fldr, "inter_*"))

        def get_vr_and_check_locpot(fldr):
            vr_file = os.path.join(fldr, "vasprun.xml")
            if not (os.path.exists(vr_file) or os.path.exists(vr_file + ".gz")):
                logger.warning("{} doesn't exit".format(vr_file))
                error_msg = ": Failure, vasprun.xml doesn't exist."
                return (None, error_msg)  # Further processing is not useful

            try:
                vr = get_vasprun(vr_file, parse_potcar_file=False)
            except:
                logger.warning("Couldn't parse {}".format(vr_file))
                error_msg = ": Failure, couldn't parse vasprun.xml file."
                return (None, error_msg)

            if not vr.converged:
                logger.warning("Vasp calculation at {} not converged".format(fldr))
                error_msg = ": Failure, Vasp calculation not converged."
                return (None, error_msg)  # Further processing is not useful

            # Check if locpot exists
            locpot_file = os.path.join(fldr, "LOCPOT")
            if not (os.path.exists(locpot_file) or os.path.exists(locpot_file + ".gz")):
                logger.warning("{} doesn't exit".format(locpot_file))
                error_msg = ": Failure, LOCPOT doesn't exist"
                return (None, error_msg)  # Further processing is not useful

            return (vr, None)

        def get_encut_from_potcar(fldr):
            potcar_file = os.path.join(fldr, "POTCAR")
            if not os.path.exists(potcar_file):
                logger.warning("Not POTCAR in {} to parse ENCUT".format(fldr))
                error_msg = ": Failure, No POTCAR file."
                return (None, error_msg)  # Further processing is not useful

            try:
                potcar = Potcar.from_file(potcar_file)
            except:
                logger.warning("Couldn't parse {}".format(potcar_file))
                error_msg = ": Failure, couldn't read POTCAR file."
                return (None, error_msg)

            encut = max(ptcr_sngl.enmax for ptcr_sngl in potcar)
            return (encut, None)

        # get bulk entry information first
        fldr = os.path.join(self._root_fldr, "bulk")
        vr, error_msg = get_vr_and_check_locpot(fldr)
        if error_msg:
            logger.error("Abandoning parsing of the calculations")
            return {}
        bulk_energy = vr.final_energy
        bulk_sc_struct = vr.final_structure
        try:
            encut = vr.incar["ENCUT"]
        except:  # ENCUT not specified in INCAR. Read from POTCAR
            encut, error_msg = get_encut_from_potcar(fldr)
            if error_msg:
                logger.error("Abandoning parsing of the calculations")
                return {}

        trans_dict = loadfn(os.path.join(fldr, "transformation.json"), cls=MontyDecoder)
        supercell_size = trans_dict["supercell"]

        bulk_file_path = fldr
        bulk_entry = ComputedStructureEntry(
            bulk_sc_struct,
            bulk_energy,
            data={"bulk_path": bulk_file_path, "encut": encut, "supercell_size": supercell_size},
        )

        # get defect entry information
        for fldr in subfolders:
            fldr_name = os.path.split(fldr)[1]
            chrg_fldrs = glob.glob(os.path.join(fldr, "charge*"))
            for chrg_fldr in chrg_fldrs:
                trans_dict = loadfn(
                    os.path.join(chrg_fldr, "transformation.json"), cls=MontyDecoder
                )
                chrg = trans_dict["charge"]
                vr, error_msg = get_vr_and_check_locpot(chrg_fldr)
                if error_msg:
                    logger.warning("Parsing the rest of the calculations")
                    continue
                if (
                    "substitution_specie" in trans_dict
                    and trans_dict["substitution_specie"] not in bulk_sc_struct.symbol_set
                ):
                    self._substitution_species.add(trans_dict["substitution_specie"])
                elif (
                    "inter" in trans_dict["defect_type"]
                    and trans_dict["defect_site"].specie.symbol not in bulk_sc_struct.symbol_set
                ):
                    # added because extrinsic interstitials don't have
                    # "substitution_specie" character...
                    trans_dict["substitution_specie"] = trans_dict["defect_site"].specie.symbol
                    self._substitution_species.add(trans_dict["defect_site"].specie.symbol)

                defect_type = trans_dict.get("defect_type", None)
                energy = vr.final_energy
                try:
                    encut = vr.incar["ENCUT"]
                except:  # ENCUT not specified in INCAR. Read from POTCAR
                    encut, error_msg = get_encut_from_potcar(chrg_fldr)
                    if error_msg:
                        logger.warning("Not able to determine ENCUT " "in {}".format(fldr_name))
                        logger.warning("Parsing the rest of the " "calculations")
                        continue

                comp_data = {
                    "bulk_path": bulk_file_path,
                    "defect_path": chrg_fldr,
                    "encut": encut,
                    "fldr_name": fldr_name,
                    "supercell_size": supercell_size,
                }
                if "substitution_specie" in trans_dict:
                    comp_data["substitution_specie"] = trans_dict["substitution_specie"]

                # create Defect Object as dict, then load to DefectEntry object
                defect_dict = {
                    "structure": bulk_sc_struct,
                    "charge": chrg,
                    "@module": "pymatgen.analysis.defects.core",
                }
                defect_site = trans_dict["defect_supercell_site"]
                if "vac_" in defect_type:
                    defect_dict["@class"] = "Vacancy"
                elif "as_" in defect_type or "sub_" in defect_type:
                    defect_dict["@class"] = "Substitution"
                    substitution_specie = trans_dict["substitution_specie"]
                    defect_site = PeriodicSite(
                        substitution_specie,
                        defect_site.frac_coords,
                        defect_site.lattice,
                        coords_are_cartesian=False,
                    )
                elif "int_" in defect_type:
                    defect_dict["@class"] = "Interstitial"
                else:
                    raise ValueError("defect type {} not recognized...".format(defect_type))

                defect_dict.update({"defect_site": defect_site})
                defect = MontyDecoder().process_decoded(defect_dict)
                parsed_defects.append(
                    DefectEntry(defect, energy - bulk_energy, parameters=comp_data)
                )

        try:
            parsed_defects_data = {}
            parsed_defects_data["bulk_entry"] = bulk_entry
            parsed_defects_data["defects"] = parsed_defects
            return parsed_defects_data
        except:
            return {}  # Return Null dict due to failure

    def get_vbm_bandgap(self):
        """
        Returns the valence band maximum (float) of the structure with
        MP-ID mpid.

        Args:
            mpid (str): MP-ID for which the valence band maximum is to
                be fetched from the Materials Project database
        """
        logger = logging.getLogger(__name__)
        vbm, bandgap = None, None

        if self._mpid is not None:
            with MPRester(api_key=self._mapi_key) as mp:
                bs = mp.get_bandstructure_by_material_id(self._mpid)
            if bs:
                vbm = bs.get_vbm()["energy"]
                bandgap = bs.get_band_gap()["energy"]

        if vbm is None or bandgap is None:
            if self._mpid:
                logger.warning(
                    "Mpid {} was provided, but no bandstructure entry currently exists for it. "
                    "Reverting to use of bulk calculation.".format(self._mpid)
                )
            else:
                logger.warning(
                    "No mp-id provided, will fetch CBM/VBM details from the " "bulk calculation."
                )
            logger.warning(
                "This may not be appropriate if the VBM/CBM occur at reciprocal points "
                "not included in the bulk calculation."
            )
            vr = get_vasprun(
                os.path.join(self._root_fldr, "bulk", "vasprun.xml"), parse_potcar_file=False
            )
            bandgap = vr.eigenvalue_band_properties[0]
            vbm = vr.eigenvalue_band_properties[2]

        return (vbm, bandgap)

    def get_chempot_limits(self):
        """
        Returns atomic chempots from bulk_composition based on data in
        the materials project database. This is abstractly handled in the
        ChemPotAnalyzer

        Note to user: If personal phase diagram desired,
            option exists in the pycdt.core.chemical_potentials to setup,
            run and parse personal phase diagrams for purposes of chemical potentials
        """
        logger = logging.getLogger(__name__)

        if self._mpid:
            cpa = chemical_potentials.MPChemPotAnalyzer(
                mpid=self._mpid, sub_species=self._substitution_species, mapi_key=self._mapi_key
            )
        else:
            bulkvr = get_vasprun(
                os.path.join(self._root_fldr, "bulk", "vasprun.xml"), parse_potcar_file=False
            )
            if not bulkvr:
                msg = "Could not fetch computed entry for atomic chempots!"
                logger.warning(msg)
                raise ValueError(msg)
            cpa = chemical_potentials.MPChemPotAnalyzer(
                bulk_ce=bulkvr.get_computed_entry(),
                sub_species=self._substitution_species,
                mapi_key=self._mapi_key,
            )

        chem_lims = cpa.analyze_GGA_chempots()

        return chem_lims

    def parse_dielectric_calculation(self):
        """
        Parses the "vasprun.xml(.gz)" file in subdirectory "dielectric" of
        root directory root_fldr and returns the average of the trace
        of the dielectric tensor.

        Args:
            root_fldr (str):
                root directory where subdirectory "dielectric" is expected
        Returns:
            eps (float):
                average of the trace of the dielectric tensor
        """

        try:
            vr = get_vasprun(
                os.path.join(self._root_fldr, "dielectric", "vasprun.xml"), parse_potcar_file=False
            )
        except:
            logging.getLogger(__name__).warning("Parsing Dielectric calculation failed")
            return None

        eps_ion = vr.epsilon_ionic
        eps_stat = vr.epsilon_static

        eps = []
        for i in range(len(eps_ion)):
            eps.append([e[0] + e[1] for e in zip(eps_ion[i], eps_stat[i])])

        return eps

    def compile_all(self):
        """
        Run to get all post processing objects as dictionary

        note: still need to implement
            1) ability for substitutional atomic chempots
            2) incorporated charge corrections for defects
        """
        output = self.parse_defect_calculations()
        output["epsilon"] = self.parse_dielectric_calculation()
        output["mu_range"] = self.get_chempot_limits()
        vbm, gap = self.get_vbm_bandgap()
        output["vbm"] = vbm
        output["gap"] = gap

        return output

max_stars_repo_name: kavanase/DefectsWithTheBoys

max_issues_repo_name: kavanase/DefectsWithTheBoys

max_forks_repo_name: kavanase/dope


text:
#!/usr/bin/python

import rosbag
import rospy
import yaml
import numpy as np

from tf.transformations import euler_from_quaternion

topic_type_dict = {}
msg_types = ['nav_msgs/Odometry', 'sensor_msgs/Imu', 'geometry_msgs/PoseStamped',
        'quadrotor_msgs/PositionCommand', 'quadrotor_msgs/TRPYCommand',
        'quadrotor_msgs/SO3Command', 'sensor_msgs/Range',
        'geometry_msgs/PoseWithCovarianceStamped']
var_types = ['x', 'y', 'z', 'vx', 'vy', 'vz',
        'acc_x', 'acc_y', 'acc_z',
        'roll', 'pitch', 'yaw',
        'ang_vel_x', 'ang_vel_y', 'ang_vel_z']

def read_bag(bagfile):
    global inbag
    inbag = rosbag.Bag(bagfile, 'r')
    return read_topic_type()

def read_topic_type():
    info_dict = yaml.load(inbag._get_yaml_info())
    for x in info_dict['topics']:
         topic_type_dict[x['topic']] = x['type']
    return topic_type_dict

def read_msg(topics):
    data = {}
    if len(topics) > 0:
        for topic, msg, type in inbag.read_messages():
         if topics.count(topic):
            if topic_type_dict[topic] == 'nav_msgs/Odometry':
                data = update_odometry(data, topic, msg)
            elif topic_type_dict[topic] == 'sensor_msgs/Imu':
                data = update_imu(data, topic, msg)
            elif topic_type_dict[topic] == 'geometry_msgs/PoseStamped':
                data = update_pose(data, topic, msg.pose, msg.header)
            elif topic_type_dict[topic] == 'quadrotor_msgs/PositionCommand':
                data = update_pose_cmd(data, topic, msg)
            elif topic_type_dict[topic] == 'geometry_msgs/PoseWithCovarianceStamped':
                data = update_pose(data, topic, msg.pose.pose, msg.header)

    return data

def update_odometry(data, topic, msg):
       quat = [msg.pose.pose.orientation.x, msg.pose.pose.orientation.y,
               msg.pose.pose.orientation.z, msg.pose.pose.orientation.w]
       [r, p, y] = euler_from_quaternion(quat)

       if topic in data:
           data[topic]['x'] = np.append(data[topic]['x'], msg.pose.pose.position.x)
           data[topic]['y'] = np.append(data[topic]['y'], msg.pose.pose.position.y)
           data[topic]['z'] = np.append(data[topic]['z'], msg.pose.pose.position.z)
           data[topic]['vx'] = np.append(data[topic]['vx'], msg.twist.twist.linear.x)
           data[topic]['vy'] = np.append(data[topic]['vy'], msg.twist.twist.linear.y)
           data[topic]['vz'] = np.append(data[topic]['vz'], msg.twist.twist.linear.z)
           data[topic]['roll'] = np.append(data[topic]['roll'], r)
           data[topic]['pitch'] = np.append(data[topic]['pitch'], p)
           data[topic]['yaw'] = np.append(data[topic]['yaw'], y)
           data[topic]['ang_vel_x'] = np.append(data[topic]['ang_vel_x'], msg.twist.twist.angular.x)
           data[topic]['ang_vel_y'] = np.append(data[topic]['ang_vel_y'], msg.twist.twist.angular.y)
           data[topic]['ang_vel_z'] = np.append(data[topic]['ang_vel_z'], msg.twist.twist.angular.z)
           data[topic]['t'] = np.append(data[topic]['t'], msg.header.stamp.to_sec())
       else:
           data[topic] = {}
           data[topic]['x'] = np.array([msg.pose.pose.position.x])
           data[topic]['y'] = np.array([msg.pose.pose.position.y])
           data[topic]['z'] = np.array([msg.pose.pose.position.z])
           data[topic]['vx'] = np.array([msg.twist.twist.linear.x])
           data[topic]['vy'] = np.array([msg.twist.twist.linear.y])
           data[topic]['vz'] = np.array([msg.twist.twist.linear.z])
           data[topic]['ang_vel_x'] = np.array([msg.twist.twist.angular.x])
           data[topic]['ang_vel_y'] = np.array([msg.twist.twist.angular.y])
           data[topic]['ang_vel_z'] = np.array([msg.twist.twist.angular.z])
           data[topic]['roll'] = np.array([r])
           data[topic]['pitch'] = np.array([p])
           data[topic]['yaw'] = np.array([y])
           data[topic]['t'] = np.array([msg.header.stamp.to_sec()])
       return data

def update_pose(data, topic, msg, header):
       quat = [msg.orientation.x, msg.orientation.y,
               msg.orientation.z, msg.orientation.w]
       [r, p, y] = euler_from_quaternion(quat)

       if topic in data:
           data[topic]['x'] = np.append(data[topic]['x'], msg.position.x)
           data[topic]['y'] = np.append(data[topic]['y'], msg.position.y)
           data[topic]['z'] = np.append(data[topic]['z'], msg.position.z)
           data[topic]['roll'] = np.append(data[topic]['roll'], r)
           data[topic]['pitch'] = np.append(data[topic]['pitch'], p)
           data[topic]['yaw'] = np.append(data[topic]['yaw'], y)
           data[topic]['t'] = np.append(data[topic]['t'], header.stamp.to_sec())
       else:
           data[topic] = {}
           data[topic]['x'] = np.array([msg.position.x])
           data[topic]['y'] = np.array([msg.position.y])
           data[topic]['z'] = np.array([msg.position.z])
           data[topic]['roll'] = np.array([r])
           data[topic]['pitch'] = np.array([p])
           data[topic]['yaw'] = np.array([y])
           data[topic]['t'] = np.array([header.stamp.to_sec()])

       return data

def update_imu(data, topic, msg):
       quat = [msg.orientation.x, msg.orientation.y,
               msg.orientation.z, msg.orientation.w]
       [r, p, y] = euler_from_quaternion(quat)

       if topic in data:
           data[topic]['acc_x'] = np.append(data[topic]['acc_x'], msg.linear_acceleration.x)
           data[topic]['acc_y'] = np.append(data[topic]['acc_y'], msg.linear_acceleration.y)
           data[topic]['acc_z'] = np.append(data[topic]['acc_z'], msg.linear_acceleration.z)
           data[topic]['ang_vel_x'] = np.append(data[topic]['ang_vel_x'], msg.angular_velocity.x)
           data[topic]['ang_vel_y'] = np.append(data[topic]['ang_vel_y'], msg.angular_velocity.y)
           data[topic]['ang_vel_z'] = np.append(data[topic]['ang_vel_z'], msg.angular_velocity.z)
           data[topic]['roll'] = np.append(data[topic]['roll'], r)
           data[topic]['pitch'] = np.append(data[topic]['pitch'], p)
           data[topic]['yaw'] = np.append(data[topic]['yaw'], y)
           data[topic]['t'] = np.append(data[topic]['t'], msg.header.stamp.to_sec())
       else:
           data[topic] = {}
           data[topic]['acc_x'] = np.array([msg.linear_acceleration.x])
           data[topic]['acc_y'] = np.array([msg.linear_acceleration.y])
           data[topic]['acc_z'] = np.array([msg.linear_acceleration.z])
           data[topic]['ang_vel_x'] = np.array([msg.angular_velocity.x])
           data[topic]['ang_vel_y'] = np.array([msg.angular_velocity.y])
           data[topic]['ang_vel_z'] = np.array([msg.angular_velocity.z])
           data[topic]['roll'] = np.array([r])
           data[topic]['pitch'] = np.array([p])
           data[topic]['yaw'] = np.array([y])
           data[topic]['t'] = np.array([msg.header.stamp.to_sec()])
       return data

def update_pose_cmd(data, topic, msg):
       if topic in data:
           data[topic]['x'] = np.append(data[topic]['x'], msg.position.x)
           data[topic]['y'] = np.append(data[topic]['y'], msg.position.y)
           data[topic]['z'] = np.append(data[topic]['z'], msg.position.z)
           data[topic]['vx'] = np.append(data[topic]['vx'], msg.velocity.x)
           data[topic]['vy'] = np.append(data[topic]['vy'], msg.velocity.y)
           data[topic]['vz'] = np.append(data[topic]['vz'], msg.velocity.z)
           data[topic]['acc_x'] = np.append(data[topic]['acc_x'], msg.acceleration.x)
           data[topic]['acc_y'] = np.append(data[topic]['acc_y'], msg.acceleration.y)
           data[topic]['acc_z'] = np.append(data[topic]['acc_z'], msg.acceleration.z)
           data[topic]['yaw'] = np.append(data[topic]['yaw'], msg.yaw)
           data[topic]['t'] = np.append(data[topic]['t'], msg.header.stamp.to_sec())
       else:
           data[topic] = {}
           data[topic]['x'] = np.array([msg.position.x])
           data[topic]['y'] = np.array([msg.position.y])
           data[topic]['z'] = np.array([msg.position.z])
           data[topic]['vx'] = np.array([msg.velocity.x])
           data[topic]['vy'] = np.array([msg.velocity.y])
           data[topic]['vz'] = np.array([msg.velocity.z])
           data[topic]['acc_x'] = np.array([msg.acceleration.x])
           data[topic]['acc_y'] = np.array([msg.acceleration.y])
           data[topic]['acc_z'] = np.array([msg.acceleration.z])
           data[topic]['yaw'] = np.array([msg.yaw])
           data[topic]['t'] = np.array([msg.header.stamp.to_sec()])
       return data

def update_trpy_cmd(data, topic, msg):
       if topic in data:
           data[topic]['roll'] = np.append(data[topic]['roll'], msg.roll)
           data[topic]['pitch'] = np.append(data[topic]['pitch'], msg.pitch)
           data[topic]['yaw'] = np.append(data[topic]['yaw'], msg.yaw)
           data[topic]['t'] = np.append(data[topic]['t'], msg.header.stamp.to_sec())
       else:
           data[topic] = {}
           data[topic]['roll'] = np.array([msg.roll])
           data[topic]['pitch'] = np.array([msg.pitch])
           data[topic]['yaw'] = np.array([msg.yaw])
           data[topic]['t'] = np.array([msg.header.stamp.to_sec()])
       return data

def update_so3_cmd(data, topic, msg):
       quat = [msg.orientation.x, msg.orientation.y,
            msg.orientation.z, msg.orientation.w]
       [r, p, y] = euler_from_quaternion(quat)

       if topic in data:
           data[topic]['yaw'] = np.append(data[topic]['yaw'], y)
           data[topic]['ang_vel_x'] = np.append(data[topic]['ang_vel_x'], msg.angular_velocity.x)
           data[topic]['ang_vel_y'] = np.append(data[topic]['ang_vel_y'], msg.angular_velocity.y)
           data[topic]['ang_vel_z'] = np.append(data[topic]['ang_vel_z'], msg.angular_velocity.z)
           data[topic]['t'] = np.append(data[topic]['t'], msg.header.stamp.to_sec())
           data[topic]['roll'] = np.append(data[topic]['roll'], r)
           data[topic]['pitch'] = np.append(data[topic]['pitch'], p)
           data[topic]['yaw'] = np.append(data[topic]['yaw'], y)
       else:
           data[topic] = {}
           data[topic]['yaw'] = np.array([y])
           data[topic]['ang_vel_x'] = np.array(msg.angular_velocity.x)
           data[topic]['ang_vel_y'] = np.array(msg.angular_velocity.y)
           data[topic]['ang_vel_z'] = np.array(msg.angular_velocity.z)
           data[topic]['t'] = np.array([msg.header.stamp.to_sec()])
           data[topic]['roll'] = np.array([r])
           data[topic]['pitch'] = np.array([p])
           data[topic]['yaw'] = np.array([y])
       return data
def update_range(data, topic, msg):
       if topic in data:
           data[topic]['z'] = np.append(data[topic]['z'], msg.range)
           data[topic]['t'] = np.append(data[topic]['t'], msg.header.stamp.to_sec())
       else:
           data[topic] = {}
           data[topic]['z'] = np.array([msg.range])
           data[topic]['t'] = np.array([msg.header.stamp.to_sec()])
       return data

if __name__ == "__main__":
    read_topic_type()

max_stars_repo_name: robomechanics/quad-software

max_issues_repo_name: robomechanics/rml-spirit-firmware

max_forks_repo_name: robomechanics/spirit-software