<commit_before>"""
Active region object definition. This object holds all the important information about our synthesized active region.
"""
import os
import sys
import logging
import datetime
import pickle
import glob

import numpy as np
import matplotlib.pyplot as plt
import seaborn.apionly as sns
import sunpy.map
import astropy.units as u
import h5py
import yt
import solarbextrapolation.map3dclasses
import solarbextrapolation.extrapolators

from synthesizAR.util import convert_angle_to_length,find_seed_points
from .loop import Loop


class Skeleton(object):
    """
    Construct magnetic field skeleton from HMI fits file

    Parameters
    ----------

    Examples
    --------
    """
    def __init__(self,hmi_fits_file=None,**kwargs):
        """
        Constructor

        Notes
        -----
        Right now, this class just accepts an HMI fits file. Could be adjusted to do the actual query as well.
        """
        self.logger = logging.getLogger(name=type(self).__name__)
        if hmi_fits_file is not None:
            tmp_map = sunpy.map.Map(hmi_fits_file)
            self._process_map(tmp_map,**kwargs)
        else:
            self.logger.warning('No HMI fits file supplied. A new HMI map object will not be created.')


    def _process_map(self,tmp_map,crop=None,resample=None):
        """
        Rotate, crop and resample map if needed. Can do any other needed processing here too.

        Parameters
        ----------
        map : `~sunpy.map.Map`
            Original HMI map
        crop : `tuple` `[xrange,yrange]`, optional
            The x- and y-ranges of the cropped map, both should be of type `~astropy.units.Quantity` and have the same units as `map.xrange` and `map.yrange`
        resample : `~astropy.units.Quantity`, `[new_xdim,new_ydim]`, optional
            The new x- and y-dimensions of the resampled map, should have the same units as `map.dimensions.x` and `map.dimensions.y`
        """

        tmp_map = tmp_map.rotate()
        if crop is not None:
            tmp_map = tmp_map.submap(crop[0],crop[1])
        if resample is not None:
            tmp_map = tmp_map.resample(resample,method='linear')

        self.hmi_map = tmp_map


    def save_field(self,savedir=None):
        """
        Save the components of the field object to be reloaded later.
        """

        top_dir = os.path.join(savedir,'synthesizAR-field-save_' + datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))
        if not os.path.exists(top_dir):
            os.makedirs(top_dir)
        #loops
        if not os.path.exists(os.path.join(top_dir,'loops')):
            os.makedirs(os.path.join(top_dir,'loops'))
        for l in self.loops:
            with open(os.path.join(top_dir,'loops',l.name+'.pickle'),'wb') as f:
                pickle.dump(l,f)
        #streamlines
        with open(os.path.join(top_dir,'streamlines.pickle'),'wb') as f:
            pickle.dump(self.streamlines,f)
        #sunpy maps
        with open(os.path.join(top_dir,'hmi_map.pickle'),'wb') as f:
            pickle.dump(self.hmi_map,f)
        with open(os.path.join(top_dir,'clipped_hmi_map.pickle'),'wb') as f:
            pickle.dump(self.clipped_hmi_map,f)
        #3d extrapolated field
        cg = self.extrapolated_3d_field.covering_grid(level=0,
                left_edge=self.extrapolated_3d_field.domain_left_edge,
                dims=self.extrapolated_3d_field.domain_dimensions)
        cg.save_as_dataset(filename=os.path.join(top_dir,
                                                'extrapolated_3d_field.h5'),
                            fields=['Bx','By','Bz'])


    @classmethod
    def restore_field(cls,savedir):
        """
        Restore the field from a set of serialized files
        """
        #loops
        loop_files = glob.glob(os.path.join(savedir,'loops','*'))
        loop_files = sorted([lf.split('/')[-1] for lf in loop_files],
                            key=lambda l: int(l.split('.')[0][4:]))
        loops = []
        for lf in loop_files:
            with open(os.path.join(savedir,'loops',lf),'rb') as f:
                 loops.append(pickle.load(f))
        #streamlines
        with open(os.path.join(savedir,'streamlines.pickle'),'rb') as f:
            streamlines = pickle.load(f)
        #sunpy maps
        with open(os.path.join(savedir,'hmi_map.pickle'),'rb') as f:
            hmi_map = pickle.load(f)
        with open(os.path.join(savedir,'clipped_hmi_map.pickle'),'rb') as f:
            clipped_hmi_map = pickle.load(f)
        #3d extapolated fields
        extrapolated_3d_field = yt.load(os.path.join(savedir,
                                                'extrapolated_3d_field.h5'))
        field = cls()
        field.loops = loops
        field.hmi_map = hmi_map
        field.clipped_hmi_map = clipped_hmi_map
        field.streamlines = streamlines
        field.extrapolated_3d_field = extrapolated_3d_field

        return field


    def _convert_angle_to_length(self,angle_or_length,working_units=u.meter):
        """
        Recast the `synthesizAR.util.convert_angle_to_length` to automatically use the supplied HMI map.
        """

        result = convert_angle_to_length(self.hmi_map,angle_or_length,working_units=working_units)
        return result


    def _transform_to_yt(self,map_3d,boundary_clipping=(2,2,2)):
        """
        Reshape data structure to something yt can work with.

        Parameters
        ----------
        map_3d : `solarbextrapolation.map3dclasses.Map3D`
            Result from the field extrapolation routine
        boundary_clipping : `tuple`, optional
            The extrapolated volume has a layer of ghost cells in each dimension. This tuple of (nx,ny,nz) tells how many cells we need to contract the volume and map in each direction.
        """

        #reshape the magnetic field data
        _tmp = map_3d.data[boundary_clipping[0]:-boundary_clipping[0], boundary_clipping[1]:-boundary_clipping[1], boundary_clipping[2]:-boundary_clipping[2],:]
        #some annoying and cryptic translation between yt and SunPy
        data = dict(
                    Bx=(np.swapaxes(_tmp[:,:,:,1],0,1),"T"),
                    By=(np.swapaxes(_tmp[:,:,:,0],0,1),"T"),
                    Bz=(np.swapaxes(_tmp[:,:,:,2],0,1),"T"))

        #trim the boundary hmi map appropriately
        self.clipped_hmi_map = self.hmi_map.submap(
                                self.hmi_map.xrange+self.hmi_map.scale.x*u.Quantity([boundary_clipping[0]*u.pixel,-boundary_clipping[0]*u.pixel]),
                                self.hmi_map.yrange+self.hmi_map.scale.y*u.Quantity([boundary_clipping[1]*u.pixel,-boundary_clipping[1]*u.pixel]))

        #create the bounding box
        bbox = np.array([
            self._convert_angle_to_length(self.clipped_hmi_map.xrange).value,
            self._convert_angle_to_length(self.clipped_hmi_map.yrange).value,
            self._convert_angle_to_length(map_3d.zrange+map_3d.scale.z*u.Quantity([boundary_clipping[2]*u.pixel,-boundary_clipping[2]*u.pixel])).value])

        #assemble the dataset
        self.extrapolated_3d_field = yt.load_uniform_grid(data,
            data['Bx'][0].shape, bbox=bbox, length_unit='cm',
            geometry=('cartesian',('x','y','z')))


    def _filter_streamlines(self,streamline,close_threshold=0.05,loop_length_range=[2.e+9,5.e+10]):
        """
        Check extracted loop to make sure it fits given criteria. Return
        True if it passes.

        Parameters
        ----------
        streamline : yt streamline object
        close_threshold : `float`
            percentage of domain width allowed between loop endpoints
        loop_length_range : `tuple`
            minimum and maximum allowed loop lengths (in centimeters)
        """

        streamline = streamline[np.all(streamline != 0.0, axis=1)]<commit_msg>better yt units<commit_after>"""
Active region object definition. This object holds all the important information about our synthesized active region.
"""
import os
import sys
import logging
import datetime
import pickle
import glob

import numpy as np
import matplotlib.pyplot as plt
import seaborn.apionly as sns
import sunpy.map
import astropy.units as u
import h5py
import yt
import solarbextrapolation.map3dclasses
import solarbextrapolation.extrapolators

from synthesizAR.util import convert_angle_to_length,find_seed_points
from .loop import Loop


class Skeleton(object):
    """
    Construct magnetic field skeleton from HMI fits file

    Parameters
    ----------

    Examples
    --------
    """
    def __init__(self,hmi_fits_file=None,**kwargs):
        """
        Constructor

        Notes
        -----
        Right now, this class just accepts an HMI fits file. Could be adjusted to do the actual query as well.
        """
        self.logger = logging.getLogger(name=type(self).__name__)
        if hmi_fits_file is not None:
            tmp_map = sunpy.map.Map(hmi_fits_file)
            self._process_map(tmp_map,**kwargs)
        else:
            self.logger.warning('No HMI fits file supplied. A new HMI map object will not be created.')


    def _process_map(self,tmp_map,crop=None,resample=None):
        """
        Rotate, crop and resample map if needed. Can do any other needed processing here too.

        Parameters
        ----------
        map : `~sunpy.map.Map`
            Original HMI map
        crop : `tuple` `[xrange,yrange]`, optional
            The x- and y-ranges of the cropped map, both should be of type `~astropy.units.Quantity` and have the same units as `map.xrange` and `map.yrange`
        resample : `~astropy.units.Quantity`, `[new_xdim,new_ydim]`, optional
            The new x- and y-dimensions of the resampled map, should have the same units as `map.dimensions.x` and `map.dimensions.y`
        """

        tmp_map = tmp_map.rotate()
        if crop is not None:
            tmp_map = tmp_map.submap(crop[0],crop[1])
        if resample is not None:
            tmp_map = tmp_map.resample(resample,method='linear')

        self.hmi_map = tmp_map


    def save_field(self,savedir=None):
        """
        Save the components of the field object to be reloaded later.
        """

        top_dir = os.path.join(savedir,'synthesizAR-field-save_' + datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))
        if not os.path.exists(top_dir):
            os.makedirs(top_dir)
        #loops
        if not os.path.exists(os.path.join(top_dir,'loops')):
            os.makedirs(os.path.join(top_dir,'loops'))
        for l in self.loops:
            with open(os.path.join(top_dir,'loops',l.name+'.pickle'),'wb') as f:
                pickle.dump(l,f)
        #streamlines
        with open(os.path.join(top_dir,'streamlines.pickle'),'wb') as f:
            pickle.dump(self.streamlines,f)
        #sunpy maps
        with open(os.path.join(top_dir,'hmi_map.pickle'),'wb') as f:
            pickle.dump(self.hmi_map,f)
        with open(os.path.join(top_dir,'clipped_hmi_map.pickle'),'wb') as f:
            pickle.dump(self.clipped_hmi_map,f)
        #3d extrapolated field
        cg = self.extrapolated_3d_field.covering_grid(level=0,
                left_edge=self.extrapolated_3d_field.domain_left_edge,
                dims=self.extrapolated_3d_field.domain_dimensions)
        cg.save_as_dataset(filename=os.path.join(top_dir,
                                                'extrapolated_3d_field.h5'),
                            fields=['Bx','By','Bz'])


    @classmethod
    def restore_field(cls,savedir):
        """
        Restore the field from a set of serialized files
        """
        #loops
        loop_files = glob.glob(os.path.join(savedir,'loops','*'))
        loop_files = sorted([lf.split('/')[-1] for lf in loop_files],
                            key=lambda l: int(l.split('.')[0][4:]))
        loops = []
        for lf in loop_files:
            with open(os.path.join(savedir,'loops',lf),'rb') as f:
                 loops.append(pickle.load(f))
        #streamlines
        with open(os.path.join(savedir,'streamlines.pickle'),'rb') as f:
            streamlines = pickle.load(f)
        #sunpy maps
        with open(os.path.join(savedir,'hmi_map.pickle'),'rb') as f:
            hmi_map = pickle.load(f)
        with open(os.path.join(savedir,'clipped_hmi_map.pickle'),'rb') as f:
            clipped_hmi_map = pickle.load(f)
        #3d extapolated fields
        extrapolated_3d_field = yt.load(os.path.join(savedir,
                                                'extrapolated_3d_field.h5'))
        field = cls()
        field.loops = loops
        field.hmi_map = hmi_map
        field.clipped_hmi_map = clipped_hmi_map
        field.streamlines = streamlines
        field.extrapolated_3d_field = extrapolated_3d_field

        return field


    def _convert_angle_to_length(self,angle_or_length,working_units=u.meter):
        """
        Recast the `synthesizAR.util.convert_angle_to_length` to automatically use the supplied HMI map.
        """

        result = convert_angle_to_length(self.hmi_map,angle_or_length,working_units=working_units)
        return result


    def _transform_to_yt(self,map_3d,boundary_clipping=(2,2,2)):
        """
        Reshape data structure to something yt can work with.

        Parameters
        ----------
        map_3d : `solarbextrapolation.map3dclasses.Map3D`
            Result from the field extrapolation routine
        boundary_clipping : `tuple`, optional
            The extrapolated volume has a layer of ghost cells in each dimension. This tuple of (nx,ny,nz) tells how many cells we need to contract the volume and map in each direction.
        """

        #reshape the magnetic field data
        _tmp = map_3d.data[boundary_clipping[0]:-boundary_clipping[0], boundary_clipping[1]:-boundary_clipping[1], boundary_clipping[2]:-boundary_clipping[2],:]
        #some annoying and cryptic translation between yt and SunPy
        data = dict(
                    Bx=(np.swapaxes(_tmp[:,:,:,1],0,1),yt.units.tesla),
                    By=(np.swapaxes(_tmp[:,:,:,0],0,1),yt.units.tesla),
                    Bz=(np.swapaxes(_tmp[:,:,:,2],0,1),yt.units.tesla))

        #trim the boundary hmi map appropriately
        self.clipped_hmi_map = self.hmi_map.submap(
                                self.hmi_map.xrange+self.hmi_map.scale.x*u.Quantity([boundary_clipping[0]*u.pixel,-boundary_clipping[0]*u.pixel]),
                                self.hmi_map.yrange+self.hmi_map.scale.y*u.Quantity([boundary_clipping[1]*u.pixel,-boundary_clipping[1]*u.pixel]))

        #create the bounding box
        bbox = np.array([
            self._convert_angle_to_length(self.clipped_hmi_map.xrange).value,
            self._convert_angle_to_length(self.clipped_hmi_map.yrange).value,
            self._convert_angle_to_length(map_3d.zrange+map_3d.scale.z*u.Quantity([boundary_clipping[2]*u.pixel,-boundary_clipping[2]*u.pixel])).value])

        #assemble the dataset
        self.extrapolated_3d_field = yt.load_uniform_grid(data,
            data['Bx'][0].shape, bbox=bbox, length_unit=yt.units.cm,
            geometry=('cartesian',('x','y','z')))


    def _filter_streamlines(self,streamline,close_threshold=0.05,loop_length_range=[2.e+9,5.e+10]):
        """
        Check extracted loop to make sure it fits given criteria. Return
        True if it passes.

        Parameters
        ----------
        streamline : yt streamline object
        close_threshold : `float`
            percentage of domain width allowed between loop endpoints
        loop_length_range : `tuple`
            minimum and maximum allowed loop lengths (in centimeters)
        """

        streamline = streamline[np.all(streamline != 0.0, axis=1)]
{'fraction_non_alphanumeric': 0.11494982852788009, 'fraction_numerical': 0.008129048647275499, 'mean_word_length': 2.4003454977326713, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2601764', 'n_tokens_mistral': 5079, 'n_tokens_neox': 4917, 'n_words': 1037}