<commit_before>#!/usr/bin/env python

"""
This module implements specific error handlers for VASP runs. These handlers
tries to detect common errors in vasp runs and attempt to fix them on the fly
by modifying the input files.
"""

from __future__ import division

__author__ = "Shyue Ping Ong"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyuep@gmail.com"
__status__ = "Beta"
__date__ = "2/4/13"

import os
import logging
import tarfile
import time
import glob
import operator

from custodian.custodian import ErrorHandler
from pymatgen.io.vaspio.vasp_input import Poscar, VaspInput
from pymatgen.transformations.standard_transformations import \
    PerturbStructureTransformation
from pymatgen.serializers.json_coders import MSONable

from pymatgen.io.vaspio.vasp_output import Vasprun
from custodian.ansible.intepreter import Modder
from custodian.ansible.actions import FileActions, DictActions

class VaspErrorHandler(ErrorHandler, MSONable):

    error_msgs = {
        "tet": ["Tetrahedron method fails for NKPT<4",
                "Fatal error detecting k-mesh",
                "Fatal error: unable to match k-point",
                "Routine TETIRR needs special values"],
        "inv_rot_mat": ["inverse of rotation matrix was not found (increase "
                        "SYMPREC)"],
        "brmix": ["BRMIX: very serious problems"],
        "rspher": ["ERROR RSPHER"],
        "subspacematrix": ["WARNING: Sub-Space-Matrix is not hermitian in DAV"],
        "tetirr": ["Routine TETIRR needs special values"],
        "incorrect_shift": ["Could not get correct shifts"],
        "mesh_symmetry": ["Reciprocal lattice and k-lattice belong to "
                          "different class of lattices."]
    }

    def __init__(self, output_filename="vasp.out"):
        self.output_filename = output_filename

    def check(self):
        self.errors = set()
        with open(self.output_filename, "r") as f:
            for line in f:
                l = line.strip()
                for err, msgs in VaspErrorHandler.error_msgs.items():
                    for msg in msgs:
                        if l.startswith(msg):
                            self.errors.add(err)
        return len(self.errors) > 0

    def correct(self):
        backup()
        actions = []
        vi = VaspInput.from_directory(".")

        if "tet" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'ISMEAR': 0}}})
        if "inv_rot_mat" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'SYMPREC': 1e-8}}})
        if "brmix" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'IMIX': 1}}})
        if "subspacematrix" in self.errors or "rspher" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'INCAR->LREAL': False}}})
        if "tetirr" in self.errors or "incorrect_shift" in self.errors:
            actions.append({'dict': 'KPOINTS',
                            'action': {'_set': {'style': "Gamma"}}})
        if "mesh_symmetry" in self.errors:
            m = reduce(operator.mul, vi["KPOINTS"].kpts[0])
            m = max(int(round(m ** (1 / 3))), 1)
            if vi["KPOINTS"].style.lower().startswith("m"):
                m += m % 2
            actions.append({'dict': 'KPOINTS',
                            'action': {'_set': {'kpoints': [[m] * 3]}}})
        m = Modder()
        modified = []
        for a in actions:
            modified.append(a["dict"])
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        for f in modified:
            vi[f].write_file(f)
        return {"errors": list(self.errors), "actions": actions}

    @property
    def is_monitor(self):
        return False

    def __str__(self):
        return "Vasp error"

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return VaspErrorHandler(d["output_filename"])

class DentetErrorHandler(ErrorHandler):

    def __init__(self, output_file='vasp.out'):
        """
        Detects an error when the output file has not been updated
        in timeout seconds. Perturbs structure and restarts
        """
        self.output_filename = output_file

    def check(self):
        with open(self.output_filename, "r") as f:
            for line in f:
                l = line.strip()
                if l.find("DENTET") != -1:
                    return True
        return False

    def correct(self):
        backup()
        actions = []
        vi = VaspInput.from_directory(".")
        actions.append({'dict': 'INCAR',
                        'action': {'_set': {'ISMEAR': 0}}})
        m = Modder()
        modified = []
        for a in actions:
            modified.append(a["dict"])
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        for f in modified:
            vi[f].write_file(f)
        return {"errors": ["dentet"], "actions": actions}

    @property
    def is_monitor(self):
        return True

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return DentetErrorHandler(d["output_filename"])

class UnconvergedErrorHandler(ErrorHandler, MSONable):
    """
    Check if a run is converged
    """
    def __init__(self, output_filename="vasprun.xml"):
        self.output_filename = output_filename

    def check(self):
        try:
            v = Vasprun(self.output_filename)
            if not v.converged:
                return True
        except:
            return True

    def correct(self):
        backup()
        actions = [{'file': 'CONTCAR',
                    'action': {'_file_copy': {'dest': 'POSCAR'}}},
                   {'dict': 'INCAR',
                    'action': {'_set': {'ISTART': 1}}}]
        vi = VaspInput.from_directory(".")
        m = Modder(actions=[DictActions, FileActions])
        for a in actions:
            if "dict" in a:
                vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
            elif "file" in a:
                m.modify(a["action"], a["file"])
        vi["INCAR"].write_file("INCAR")

        return {"errors": ["Unconverged"], "actions": actions}

    def __str__(self):
        return "Run unconverged."

    @property
    def is_monitor(self):
        return False

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return UnconvergedErrorHandler(d["output_filename"])

class PoscarErrorHandler(ErrorHandler, MSONable):

    def __init__(self, output_filename="vasp.out"):
        self.output_filename = output_filename

    def check(self):
        with open(self.output_filename, "r") as f:
            output = f.read()
            for line in output.split("\n"):
                l = line.strip()
                if l.startswith("Found some non-integer element in rotation "
                                "matrix"):
                    return True
        return False

    def correct(self):
        backup()
        p = Poscar.from_file("POSCAR")
        s = p.structure
        trans = PerturbStructureTransformation(0.05)
        new_s = trans.apply_transformation(s)
        actions = [{'dict': 'POSCAR',
                    'action': {'_set': {'structure': new_s.to_dict}}}]
        m = Modder()
        vi = VaspInput.from_directory(".")
        for a in actions:
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        vi["POSCAR"].write_file("POSCAR")

        return {"errors": ["Rotation matrix"],
                "actions": actions}

    @property
    def is_monitor(self):
        return False

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return PoscarErrorHandler(d["output_filename"])

class FrozenJobErrorHandler(ErrorHandler):

    def __init__(self, output_filename='vasp.out', timeout=3600):
        """
        Detects an error when the output file has not been updated
        in timeout seconds. Perturbs structure and restarts
        """
        self.output_filename = output_filename
        self.timeout = timeout

    def check(self):
        st = os.stat(self.output_filename)
        if time.time() - st.st_mtime > self.timeout:
            return True

    def correct(self):
        backup()
        p = Poscar.from_file("POSCAR")
        s = p.structure
        trans = PerturbStructureTransformation(0.05)
        new_s = trans.apply_transformation(s)
        actions = [{'dict': 'POSCAR',
                    'action': {'_set': {'structure': new_s.to_dict}}}]
        m = Modder()
        vi = VaspInput.from_directory(".")
        for a in actions:
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        vi["POSCAR"].write_file("POSCAR")

        return {"errors": ["Frozen job"],
                "actions": actions}

    @property
    def is_monitor(self):
        return True

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename,
                "timeout": self.timeout}

    @staticmethod
    def from_dict(d):
        return FrozenJobErrorHandler(d["output_filename"],
                                     timeout=d["timeout"])

def backup():
    error_num = 0
    for f in glob.glob("error.*.tar.gz"):
        toks = f.split(".")
        error_num = max(error_num, int(toks[1]))
    filename = "error.{}.tar.gz".format(error_num + 1)
    logging.info("Backing up run to {}.".format(filename))
    tar = tarfile.open(filename, "w:gz")
    for f in os.listdir("."):
        if not (f.startswith("error") and f.endswith(".tar.gz")):
            tar.add(f)
    tar.close()
<commit_msg>Fix bug in custodian for incorrect shift.<commit_after>#!/usr/bin/env python

"""
This module implements specific error handlers for VASP runs. These handlers
tries to detect common errors in vasp runs and attempt to fix them on the fly
by modifying the input files.
"""

from __future__ import division

__author__ = "Shyue Ping Ong"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyuep@gmail.com"
__status__ = "Beta"
__date__ = "2/4/13"

import os
import logging
import tarfile
import time
import glob
import operator

from custodian.custodian import ErrorHandler
from pymatgen.io.vaspio.vasp_input import Poscar, VaspInput
from pymatgen.transformations.standard_transformations import \
    PerturbStructureTransformation
from pymatgen.serializers.json_coders import MSONable

from pymatgen.io.vaspio.vasp_output import Vasprun
from custodian.ansible.intepreter import Modder
from custodian.ansible.actions import FileActions, DictActions

class VaspErrorHandler(ErrorHandler, MSONable):

    error_msgs = {
        "tet": ["Tetrahedron method fails for NKPT<4",
                "Fatal error detecting k-mesh",
                "Fatal error: unable to match k-point",
                "Routine TETIRR needs special values"],
        "inv_rot_mat": ["inverse of rotation matrix was not found (increase "
                        "SYMPREC)"],
        "brmix": ["BRMIX: very serious problems"],
        "rspher": ["ERROR RSPHER"],
        "subspacematrix": ["WARNING: Sub-Space-Matrix is not hermitian in DAV"],
        "tetirr": ["Routine TETIRR needs special values"],
        "incorrect_shift": ["Could not get correct shifts"],
        "mesh_symmetry": ["Reciprocal lattice and k-lattice belong to "
                          "different class of lattices."]
    }

    def __init__(self, output_filename="vasp.out"):
        self.output_filename = output_filename

    def check(self):
        self.errors = set()
        with open(self.output_filename, "r") as f:
            for line in f:
                l = line.strip()
                for err, msgs in VaspErrorHandler.error_msgs.items():
                    for msg in msgs:
                        if l.startswith(msg):
                            self.errors.add(err)
        return len(self.errors) > 0

    def correct(self):
        backup()
        actions = []
        vi = VaspInput.from_directory(".")

        if "tet" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'ISMEAR': 0}}})
        if "inv_rot_mat" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'SYMPREC': 1e-8}}})
        if "brmix" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'IMIX': 1}}})
        if "subspacematrix" in self.errors or "rspher" in self.errors:
            actions.append({'dict': 'INCAR',
                            'action': {'_set': {'LREAL': False}}})
        if "tetirr" in self.errors or "incorrect_shift" in self.errors:
            actions.append({'dict': 'KPOINTS',
                            'action': {'_set': {'generation_style': "Gamma"}}})
        if "mesh_symmetry" in self.errors:
            m = reduce(operator.mul, vi["KPOINTS"].kpts[0])
            m = max(int(round(m ** (1 / 3))), 1)
            if vi["KPOINTS"].style.lower().startswith("m"):
                m += m % 2
            actions.append({'dict': 'KPOINTS',
                            'action': {'_set': {'kpoints': [[m] * 3]}}})
        m = Modder()
        modified = []
        for a in actions:
            modified.append(a["dict"])
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        for f in modified:
            vi[f].write_file(f)
        return {"errors": list(self.errors), "actions": actions}

    @property
    def is_monitor(self):
        return False

    def __str__(self):
        return "Vasp error"

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return VaspErrorHandler(d["output_filename"])

class DentetErrorHandler(ErrorHandler):

    def __init__(self, output_file='vasp.out'):
        """
        Detects an error when the output file has not been updated
        in timeout seconds. Perturbs structure and restarts
        """
        self.output_filename = output_file

    def check(self):
        with open(self.output_filename, "r") as f:
            for line in f:
                l = line.strip()
                if l.find("DENTET") != -1:
                    return True
        return False

    def correct(self):
        backup()
        actions = []
        vi = VaspInput.from_directory(".")
        actions.append({'dict': 'INCAR',
                        'action': {'_set': {'ISMEAR': 0}}})
        m = Modder()
        modified = []
        for a in actions:
            modified.append(a["dict"])
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        for f in modified:
            vi[f].write_file(f)
        return {"errors": ["dentet"], "actions": actions}

    @property
    def is_monitor(self):
        return True

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return DentetErrorHandler(d["output_filename"])

class UnconvergedErrorHandler(ErrorHandler, MSONable):
    """
    Check if a run is converged
    """
    def __init__(self, output_filename="vasprun.xml"):
        self.output_filename = output_filename

    def check(self):
        try:
            v = Vasprun(self.output_filename)
            if not v.converged:
                return True
        except:
            return True

    def correct(self):
        backup()
        actions = [{'file': 'CONTCAR',
                    'action': {'_file_copy': {'dest': 'POSCAR'}}},
                   {'dict': 'INCAR',
                    'action': {'_set': {'ISTART': 1}}}]
        vi = VaspInput.from_directory(".")
        m = Modder(actions=[DictActions, FileActions])
        for a in actions:
            if "dict" in a:
                vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
            elif "file" in a:
                m.modify(a["action"], a["file"])
        vi["INCAR"].write_file("INCAR")

        return {"errors": ["Unconverged"], "actions": actions}

    def __str__(self):
        return "Run unconverged."

    @property
    def is_monitor(self):
        return False

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return UnconvergedErrorHandler(d["output_filename"])

class PoscarErrorHandler(ErrorHandler, MSONable):

    def __init__(self, output_filename="vasp.out"):
        self.output_filename = output_filename

    def check(self):
        with open(self.output_filename, "r") as f:
            output = f.read()
            for line in output.split("\n"):
                l = line.strip()
                if l.startswith("Found some non-integer element in rotation "
                                "matrix"):
                    return True
        return False

    def correct(self):
        backup()
        p = Poscar.from_file("POSCAR")
        s = p.structure
        trans = PerturbStructureTransformation(0.05)
        new_s = trans.apply_transformation(s)
        actions = [{'dict': 'POSCAR',
                    'action': {'_set': {'structure': new_s.to_dict}}}]
        m = Modder()
        vi = VaspInput.from_directory(".")
        for a in actions:
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        vi["POSCAR"].write_file("POSCAR")

        return {"errors": ["Rotation matrix"],
                "actions": actions}

    @property
    def is_monitor(self):
        return False

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @staticmethod
    def from_dict(d):
        return PoscarErrorHandler(d["output_filename"])

class FrozenJobErrorHandler(ErrorHandler):

    def __init__(self, output_filename='vasp.out', timeout=3600):
        """
        Detects an error when the output file has not been updated
        in timeout seconds. Perturbs structure and restarts
        """
        self.output_filename = output_filename
        self.timeout = timeout

    def check(self):
        st = os.stat(self.output_filename)
        if time.time() - st.st_mtime > self.timeout:
            return True

    def correct(self):
        backup()
        p = Poscar.from_file("POSCAR")
        s = p.structure
        trans = PerturbStructureTransformation(0.05)
        new_s = trans.apply_transformation(s)
        actions = [{'dict': 'POSCAR',
                    'action': {'_set': {'structure': new_s.to_dict}}}]
        m = Modder()
        vi = VaspInput.from_directory(".")
        for a in actions:
            vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
        vi["POSCAR"].write_file("POSCAR")

        return {"errors": ["Frozen job"],
                "actions": actions}

    @property
    def is_monitor(self):
        return True

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename,
                "timeout": self.timeout}

    @staticmethod
    def from_dict(d):
        return FrozenJobErrorHandler(d["output_filename"],
                                     timeout=d["timeout"])

def backup():
    error_num = 0
    for f in glob.glob("error.*.tar.gz"):
        toks = f.split(".")
        error_num = max(error_num, int(toks[1]))
    filename = "error.{}.tar.gz".format(error_num + 1)
    logging.info("Backing up run to {}.".format(filename))
    tar = tarfile.open(filename, "w:gz")
    for f in os.listdir("."):
        if not (f.startswith("error") and f.endswith(".tar.gz")):
            tar.add(f)
    tar.close()
<|endoftext|>


<commit_before>import numpy as np
import os

from astropy.table import Table
from . import utils

__all__ = ["FilterDefinition", "FilterFile", "ParamFilter"]

VEGA_FILE = os.path.join(utils.path_to_eazy_data(),
                         'alpha_lyr_stis_008.fits')

VEGA = Table.read(VEGA_FILE)
for c in VEGA.colnames:
    VEGA[c] = VEGA[c].astype(float)

class FilterDefinition:
    def __init__(self, name=None, wave=None, throughput=None, bp=None):
        """
        Bandpass object

        Parameters
        ----------
        name : str
            Label name

        wave : array
            Wavelength array, in `astropy.units.Angstrom`.

        throughput : array
            Throughput, arbitrary normalization

        bp : optional, `pysynphot.obsbandpass` object
            `pysynphot` filter bandpass

        """
        self.name = name
        self.wave = wave
        self.throughput = throughput
        self.Aflux = 1.

        # pysynphot Bandpass
        if bp is not None:
            self.wave = np.cast[np.double](bp.wave)
            self.throughput =  np.cast[np.double](bp.throughput)
            self.name = bp.name

        self.norm = 1.
        if self.throughput is not None:
            self.norm = np.trapz(self.throughput/self.wave, self.wave)

    def __repr__(self):
        return self.name.__repr__()

    def __str__(self):
        return self.name.__str__()

    def get_extinction(self, EBV=0, Rv=3.1):
        """
        Extinction factor
        """
        import astropy.units as u

        f99 = utils.GalacticExtinction(EBV=EBV, Rv=Rv)
        self.Alambda = f99(self.wave)
        self.Aflux = 10**(-0.4*self.Alambda)

    def extinction_correction(self, EBV, Rv=3.1, mag=True, source_lam=None, source_flux=None):
        """
        Get the MW extinction correction within the filter.

        Optionally supply a source spectrum.
        """
        import astropy.units as u
        try:
            import grizli.utils_c
            interp = grizli.utils_c.interp.interp_conserve_c
        except ImportError:
            interp = utils.interp_conserve

        if self.wave is None:
            print('Filter not defined.')
            return False

        if source_flux is None:
            source_flux = self.throughput*0.+1
        else:
            source_flux = interp(self.wave, source_lam, source_flux, left=0, right=0)

        if (self.wave.min() < 910) | (self.wave.max() > 6.e4):
            Alambda = 0.
        else:
            f99 = utils.GalacticExtinction(EBV=EBV, Rv=Rv)
            Alambda = f99(self.wave)

        delta = np.trapz(self.throughput*source_flux*10**(-0.4*Alambda), self.wave) / np.trapz(self.throughput*source_flux, self.wave)

        if mag:
            return 2.5*np.log10(delta)
        else:
            return 1./delta

    @property
    def ABVega(self):
        """
        Compute AB-Vega conversion
        """
        from astropy.constants import c
        import astropy.units as u
        try:
            import grizli.utils_c
            interp = grizli.utils_c.interp.interp_conserve_c
        except ImportError:
            interp = utils.interp_conserve

        # Union of throughput and Vega spectrum arrays
        full_x = np.hstack([self.wave, VEGA['WAVELENGTH']])
        full_x = full_x[np.argsort(full_x)]

        # Vega spectrum, units of f-lambda flux density, cgs
        # Interpolate to wavelength grid, no extrapolation
        vega_full = interp(full_x, VEGA['WAVELENGTH'], VEGA['FLUX'],
                              left=0, right=0)

        thru_full = interp(full_x, self.wave, self.throughput,
                              left=0, right=0)

        # AB = 0, same units
        absp = 3631*1e-23*c.to(u.m/u.s).value*1.e10/full_x**2

        # Integrate over the bandpass, flam dlam
        num = np.trapz(vega_full*thru_full, full_x)
        den = np.trapz(absp*thru_full, full_x)

        return -2.5*np.log10(num/den)

    @property
    def pivot(self):
        """
        Pivot wavelength

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """
        integrator = np.trapz

        num = integrator(self.wave, self.wave*self.throughput)
        den = integrator(self.wave, self.throughput/self.wave)
        pivot = np.sqrt(num/den)
        return pivot

    @property
    def equivwidth(self):
        """
        Filter equivalent width

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """
        return np.trapz(self.throughput, self.wave)

    @property
    def rectwidth(self):
        """
        Filter rectangular width

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """

        rect = self.equivwidth / self.throughput.max()
        return rect

    @property
    def ctw95(self):
        """
        95% cumulative throughput width

        http://www.stsci.edu/hst/acs/analysis/bandwidths/#keywords

        """

        dl = np.diff(self.wave)
        filt = np.cumsum((self.wave*self.throughput)[1:]*dl)
        ctw95 = np.interp([0.025, 0.975], filt/filt.max(), self.wave[1:])
        return np.diff(ctw95)[0]

    def for_filter_file(self, row_str='{i:6} {wave:.5e} {thru:.5e}'):
        """
        Return a string that can be put in the EAZY filter file
        """
        header = '{0} {1} lambda_c= {2:.4e} AB-Vega= {3:.3f} w95={4:.1f}'
        N = len(self.wave)
        lines = [header.format(N, self.name.split('lambda_c')[0],
                               self.pivot, self.ABVega, self.ctw95)]

        lines += [row_str.format(i=i+1, wave=w, thru=t)
                  for i, (w, t) in enumerate(zip(self.wave, self.throughput))]

        return '\n'.join(lines)

class FilterFile:
    def __init__(self, file='FILTER.RES.latest', path='./'):
        """
        Read a EAZY filter file.

        .. plot::
            :include-source:

            import matplotlib.pyplot as plt
            from eazy.filters import FilterFile

            res = FilterFile(path=None)
            print(len(res.filters))

            bp = res[205]
            print(bp)

            fig, ax = plt.subplots(1,1,figsize=(6,4))

            ax.plot(bp.wave, bp.throughput, label=bp.name.split()[0])

            ax.set_xlabel('wavelength, Angstroms')
            ax.set_ylabel('throughput')
            ax.legend()
            ax.grid()

            fig.tight_layout(pad=0.5)

        """
        if path is None:
            file_path = os.path.join(os.getenv('EAZYCODE'), 'filters', file)
        else:
            file_path = os.path.join(path, file)

        with open(file_path, 'r') as fp:
            lines = fp.readlines()

        self.filename = file_path

        filters = []
        wave = []
        trans = []
        header = ''

        for line in lines:
            if 'lambda_c' in line:
                if len(wave) > 0:
                    # Make filter from lines already read in
                    new_filter = FilterDefinition(name=header,
                                                  wave=np.cast[float](wave),
                                            throughput=np.cast[float](trans))
                    # new_filter.name = header
                    # new_filter.wave = np.cast[float](wave)
                    # new_filter.throughput = np.cast[float](trans)
                    filters.append(new_filter)

                # Initialize filter
                header = ' '.join(line.split()[1:])
                wave = []
                trans = []
            else:
                lspl = np.cast[float](line.split())
                wave.append(lspl[1])
                trans.append(lspl[2])

        # last one
        # new_filter = FilterDefinition()
        # new_filter.name = header
        # new_filter.wave = np.cast[float](wave)
        # new_filter.throughput = np.cast[float](trans)
        new_filter = FilterDefinition(name=header,
                                      wave=np.cast[float](wave),
                                throughput=np.cast[float](trans))

        filters.append(new_filter)

        self.filters = filters

    @property
    def NFILT(self):
        """
        Number of filters in the list
        """
        return len(self.filters)

    def __getitem__(self, i1):
        """
        Return unit-indexed filter, e.g., 161 = 2mass-j
        """
        return self.filters[i1-1]

    def names(self, verbose=True):
        """
        Print the filter names.
        """
        if verbose:
            for i in range(len(self.filters)):
                print('{0:5d} {1}'.format(i+1, self.filters[i].name))
        else:
            string_list = ['{0:5d} {1}\n'.format(i+1, self.filters[i].name) for i in range(len(self.filters))]
            return string_list

    def write(self, file='xxx.res', verbose=True):
        """
        Dump the filter information to a filter file.
        """
        fp = open(file,'w')
        for filter in self.filters:
            fp.write('{0:6d} {1}\n'.format(len(filter.wave), filter.name))
            for i in range(len(filter.wave)):
                fp.write('{0:6d} {1:.5e} {2:.5e}\n'.format(i+1, filter.wave[i], filter.throughput[i]))

        fp.close()

        string_list = self.names(verbose=False)
        fp = open(file+'.info', 'w')
        fp.writelines(string_list)
        fp.close()

        if verbose:
            print('Wrote <{0}[.info]>'.format(file))

    def search(self, search_string, case=False, verbose=True):
        """
        Search filter names for ``search_string``.  If ``case`` is True, then
        match case.
        """
        import re

        if not case:
            search_string = search_string.upper()

        matched = []

        for i in range(len(self.filters)):
            filt_name = self.filters[i].name
            if not case:
                filt_name = filt_name.upper()

            if re.search(search_string, filt_name) is not None:
                if verbose:
                    print('{0:5d} {1}'.format(i+1, self.filters[i].name))
                matched.append(i)

        return np.array(matched)

class ParamFilter(FilterDefinition):
    def __init__(self, line='#  Filter #20, RES#78: COSMOS/SUBARU_filter_B.txt - lambda_c=4458.276253'):

        self.lambda_c = float(line.split('lambda_c=')[1])
        self.name = line.split()[4]
        self.fnumber = int(line.split('RES#')[1].split(':')[0])
        self.cnumber = int(line.split('Filter #')[1].split(',')[0])
<commit_msg>add switch for photon- and energy-counting filters<commit_after>import numpy as np
import os

from astropy.table import Table
from . import utils

__all__ = ["FilterDefinition", "FilterFile", "ParamFilter"]

VEGA_FILE = os.path.join(utils.path_to_eazy_data(),
                         'alpha_lyr_stis_008.fits')

VEGA = Table.read(VEGA_FILE)
for c in VEGA.colnames:
    VEGA[c] = VEGA[c].astype(float)

class FilterDefinition:
    def __init__(self, name=None, wave=None, throughput=None, bp=None, photon_counter=True):
        """
        Bandpass object

        Parameters
        ----------
        name : str
            Label name

        wave : array
            Wavelength array, in `astropy.units.Angstrom`.

        throughput : array
            Throughput, arbitrary normalization

        bp : optional, `pysynphot.obsbandpass` object
            `pysynphot` filter bandpass

        photon_counter : bool
            Filter is associated with a photon-counting detector (e.g. CCDs).
            Often FIR devices are energy-counting, e.g., MIPS, Herschel.

        """
        self.name = name
        self.wave = wave
        self._throughput = throughput
        self.Aflux = 1.
        self.photon_counter = photon_counter

        if isinstance(name, str):
            # Set energy-counting filters
            ir_filts = ['mips/24','mips/70','mips/160',
                        'scuba2/450','scuba2/850',
                        'herschel/pacs/70', 'herschel/pacs/100',
                        'herschel/pacs/160', 'herschel/spire/200',
                        'herschel/spire/350', 'herschel/spire/500']

            if name.split()[0] in ir_filts:
                self.photon_counter = False

        # pysynphot Bandpass
        if bp is not None:
            self.wave = np.cast[np.double](bp.wave)
            self._throughput =  np.cast[np.double](bp.throughput)
            self.name = bp.name

        # Set throughput accounting for photon_counter
        self.set_throughput()

    def set_throughput(self):
        """
        Set throughput accounting for `photon_counter` attribute
        """
        if self._throughput is None:
            self.throughput = None
            self.norm = 1.
        else:
            if self.photon_counter:
                self.throughput = self._throughput
            else:
                _wfact = self.wave/np.mean(self.wave)
                self.throughput =  self._throughput/_wfact

            self.norm = np.trapz(self.throughput/self.wave, self.wave)

    def __repr__(self):
        return self.name.__repr__()

    def __str__(self):
        return self.name.__str__()

    def get_extinction(self, EBV=0, Rv=3.1):
        """
        Extinction factor
        """
        import astropy.units as u

        f99 = utils.GalacticExtinction(EBV=EBV, Rv=Rv)
        self.Alambda = f99(self.wave)
        self.Aflux = 10**(-0.4*self.Alambda)

    def extinction_correction(self, EBV, Rv=3.1, mag=True, source_lam=None, source_flux=None):
        """
        Get the MW extinction correction within the filter.

        Optionally supply a source spectrum.
        """
        import astropy.units as u
        try:
            import grizli.utils_c
            interp = grizli.utils_c.interp.interp_conserve_c
        except ImportError:
            interp = utils.interp_conserve

        if self.wave is None:
            print('Filter not defined.')
            return False

        if source_flux is None:
            src = self.throughput*0.+1
        else:
            src = interp(self.wave, source_lam, source_flux, left=0, right=0)

        if (self.wave.min() < 910) | (self.wave.max() > 6.e4):
            Alambda = 0.
        else:
            f99 = utils.GalacticExtinction(EBV=EBV, Rv=Rv)
            Alambda = f99(self.wave)

        src_red = np.trapz(self.throughput*src*10**(-0.4*Alambda)/self.wave,
                           self.wave)
        src_nored = np.trapz(self.throughput*src/self.wave, self.wave)

        delta = src_red/src_nored

        if mag:
            return 2.5*np.log10(delta)
        else:
            return 1./delta

    @property
    def ABVega(self):
        """
        Compute AB-Vega conversion
        """
        from astropy.constants import c
        import astropy.units as u
        try:
            import grizli.utils_c
            interp = grizli.utils_c.interp.interp_conserve_c
        except ImportError:
            interp = utils.interp_conserve

        # Union of throughput and Vega spectrum arrays
        full_x = np.hstack([self.wave, VEGA['WAVELENGTH']])
        full_x = full_x[np.argsort(full_x)]

        # Vega spectrum, units of f-lambda flux density, cgs
        # Interpolate to wavelength grid, no extrapolation
        vega_full = interp(full_x, VEGA['WAVELENGTH'], VEGA['FLUX'],
                              left=0, right=0)

        thru_full = interp(full_x, self.wave, self.throughput,
                              left=0, right=0)

        # AB = 0, same units
        absp = 3631*1e-23*c.to(u.m/u.s).value*1.e10/full_x**2

        # Integrate over the bandpass, flam dlam
        num = np.trapz(vega_full*thru_full, full_x)
        den = np.trapz(absp*thru_full, full_x)

        return -2.5*np.log10(num/den)

    @property
    def pivot(self):
        """
        Pivot wavelength

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """
        integrator = np.trapz

        num = integrator(self.wave, self.wave*self.throughput)
        den = integrator(self.wave, self.throughput/self.wave)
        pivot = np.sqrt(num/den)
        return pivot

    @property
    def equivwidth(self):
        """
        Filter equivalent width

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """
        return np.trapz(self.throughput, self.wave)

    @property
    def rectwidth(self):
        """
        Filter rectangular width

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """

        rect = self.equivwidth / self.throughput.max()
        return rect

    @property
    def ctw95(self):
        """
        95% cumulative throughput width

        http://www.stsci.edu/hst/acs/analysis/bandwidths/#keywords

        """

        dl = np.diff(self.wave)
        filt = np.cumsum((self.wave*self.throughput)[1:]*dl)
        ctw95 = np.interp([0.025, 0.975], filt/filt.max(), self.wave[1:])
        return np.diff(ctw95)[0]

    def for_filter_file(self, row_str='{i:6} {wave:.5e} {thru:.5e}'):
        """
        Return a string that can be put in the EAZY filter file
        """
        header = '{0} {1} lambda_c= {2:.4e} AB-Vega= {3:.3f} w95={4:.1f}'
        N = len(self.wave)
        lines = [header.format(N, self.name.split('lambda_c')[0],
                               self.pivot, self.ABVega, self.ctw95)]

        lines += [row_str.format(i=i+1, wave=w, thru=t)
                  for i, (w, t) in enumerate(zip(self.wave, self.throughput))]

        return '\n'.join(lines)

class FilterFile:
    def __init__(self, file='FILTER.RES.latest', path='./'):
        """
        Read a EAZY filter file.

        .. plot::
            :include-source:

            import matplotlib.pyplot as plt
            from eazy.filters import FilterFile

            res = FilterFile(path=None)
            print(len(res.filters))

            bp = res[205]
            print(bp)

            fig, ax = plt.subplots(1,1,figsize=(6,4))

            ax.plot(bp.wave, bp.throughput, label=bp.name.split()[0])

            ax.set_xlabel('wavelength, Angstroms')
            ax.set_ylabel('throughput')
            ax.legend()
            ax.grid()

            fig.tight_layout(pad=0.5)

        """
        if path is None:
            file_path = os.path.join(os.getenv('EAZYCODE'), 'filters', file)
        else:
            file_path = os.path.join(path, file)

        with open(file_path, 'r') as fp:
            lines = fp.readlines()

        self.filename = file_path

        filters = []
        wave = []
        trans = []
        header = ''

        for line in lines:
            if 'lambda_c' in line:
                if len(wave) > 0:
                    # Make filter from lines already read in
                    new_filter = FilterDefinition(name=header,
                                                  wave=np.cast[float](wave),
                                            throughput=np.cast[float](trans))
                    # new_filter.name = header
                    # new_filter.wave = np.cast[float](wave)
                    # new_filter.throughput = np.cast[float](trans)
                    filters.append(new_filter)

                # Initialize filter
                header = ' '.join(line.split()[1:])
                wave = []
                trans = []
            else:
                lspl = np.cast[float](line.split())
                wave.append(lspl[1])
                trans.append(lspl[2])

        # last one
        # new_filter = FilterDefinition()
        # new_filter.name = header
        # new_filter.wave = np.cast[float](wave)
        # new_filter.throughput = np.cast[float](trans)
        new_filter = FilterDefinition(name=header,
                                      wave=np.cast[float](wave),
                                throughput=np.cast[float](trans))

        filters.append(new_filter)

        self.filters = filters

    @property
    def NFILT(self):
        """
        Number of filters in the list
        """
        return len(self.filters)

    def __getitem__(self, i1):
        """
        Return unit-indexed filter, e.g., 161 = 2mass-j
        """
        return self.filters[i1-1]

    def names(self, verbose=True):
        """
        Print the filter names.
        """
        if verbose:
            for i in range(len(self.filters)):
                print('{0:5d} {1}'.format(i+1, self.filters[i].name))
        else:
            string_list = ['{0:5d} {1}\n'.format(i+1, self.filters[i].name) for i in range(len(self.filters))]
            return string_list

    def write(self, file='xxx.res', verbose=True):
        """
        Dump the filter information to a filter file.
        """
        fp = open(file,'w')
        for filter in self.filters:
            fp.write('{0:6d} {1}\n'.format(len(filter.wave), filter.name))
            for i in range(len(filter.wave)):
                fp.write('{0:6d} {1:.5e} {2:.5e}\n'.format(i+1, filter.wave[i], filter.throughput[i]))

        fp.close()

        string_list = self.names(verbose=False)
        fp = open(file+'.info', 'w')
        fp.writelines(string_list)
        fp.close()

        if verbose:
            print('Wrote <{0}[.info]>'.format(file))

    def search(self, search_string, case=False, verbose=True):
        """
        Search filter names for ``search_string``.  If ``case`` is True, then
        match case.
        """
        import re

        if not case:
            search_string = search_string.upper()

        matched = []

        for i in range(len(self.filters)):
            filt_name = self.filters[i].name
            if not case:
                filt_name = filt_name.upper()

            if re.search(search_string, filt_name) is not None:
                if verbose:
                    print('{0:5d} {1}'.format(i+1, self.filters[i].name))
                matched.append(i)

        return np.array(matched)

class ParamFilter(FilterDefinition):
    def __init__(self, line='#  Filter #20, RES#78: COSMOS/SUBARU_filter_B.txt - lambda_c=4458.276253'):

        self.lambda_c = float(line.split('lambda_c=')[1])
        self.name = line.split()[4]
        self.fnumber = int(line.split('RES#')[1].split(':')[0])
        self.cnumber = int(line.split('Filter #')[1].split(',')[0])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import datetime
from dateutil.relativedelta import relativedelta
import calendar
from itertools import islice
import math
from decimal import Decimal
from collections import defaultdict

class Payment(object):
    def __init__(self, date, payment, interest_payment, principal_payment, current_principal):
        self.date = date
        self.payment = payment
        self.interest_payment = interest_payment
        self.principal_payment = principal_payment
        self.current_principal = current_principal

class RegularPayment(Payment):
    def __init__(self, date, payment, interest_payment, principal_payment, current_principal):
        return super().__init__(date, payment, interest_payment, principal_payment, current_principal)

    def __str__(self):
        return '%s - %.2f - %.2f - %.2f - %.2f - regular' % (str(self.date), self.payment, self.interest_payment, self.principal_payment, self.current_principal)

class NonRegularPayment(Payment):
    def __init__(self, date, payment, interest_payment, principal_payment, current_principal):
        super().__init__(date, payment, interest_payment, principal_payment, current_principal)

    def __str__(self):
        return '%s - %.2f - %.2f - %.2f - %.2f - non regular' % (str(self.date), self.payment, self.interest_payment, self.principal_payment, self.current_principal)

class Period(object):
    def __init__(self, start, end):
        self.start = start
        self.end = end
        self.payments = []
        self.principal_paid = False

    def add_payment(self, payment):
        self.payments.append(payment)

class Result(object):
    def __init__(self):
        self.payments = defaultdict(list)
        self.year_interest_payments = defaultdict(Decimal)

class Mortgage(object):
    def __init__(self, interest, months, start_date):
        self.interest = interest
        self.interest_rate = Decimal(interest / 100)
        self.monthly_interest = Decimal(self.interest_rate / 12)

        # срок кредита
        self.total_months = months

        # дата получения кредита
        self.start_date = start_date

        self.cached_calendar = [self.start_date + relativedelta(months=m) for m in range(0, self.total_months + 1)]

        self.leap_years_cache = [366 if calendar.isleap(self.start_date.year + y) else 365 for y in range(0, int(self.total_months / 12))]

    def calc_monthly_payment(self, principal, months):
        return Decimal(round(principal * (self.monthly_interest / (1 - (1 + self.monthly_interest) ** (- months))), 2))

    def calc_interest_payment(self, pay_date, prev_pay_date, principal):
        # первый месяц в году, т.е. январь
        # в этом случае надо немного по-другому считать
        days_in_year = self.leap_years_cache[self.start_date.year - pay_date.year]
        if pay_date.month == 1 and prev_pay_date.month == 12:
            december31 = datetime.date(prev_pay_date.year, 12, 31)

            days1 = december31 - prev_pay_date
            days1 = days1.days
            p1 = Decimal(days1 / self.leap_years_cache[self.start_date.year - prev_pay_date.year])
            interest_payment1 = Decimal(round(principal * self.interest_rate * p1, 2))

            days2 = pay_date - december31
            days2 = days2.days
            p2 = Decimal(days2 / days_in_year)
            interest_payment2 = Decimal(round(principal * self.interest_rate * p2, 2))

            return interest_payment1 + interest_payment2
        else:
            d = pay_date - prev_pay_date
            days = d.days

            return Decimal(round(principal * self.interest_rate * days / days_in_year, 2))

    def is_pay_in_regular_date(self, pay):
        return pay.date.day == self.start_date.day

    def calc(self, principal, predefined_payments = {}, initial_month = None):
        current_principal = principal
        monthly_payment = self.calc_monthly_payment(current_principal, self.total_months)

        result = Result()
        cur_period = self.cached_calendar[0 if initial_month == None else initial_month]

        for m in range(1 if initial_month == None else initial_month + 1, self.total_months + 1):
            end_period = self.cached_calendar[m]
            principal_paid = False

            non_reg_payments = predefined_payments.get((end_period.year, end_period.month), [])
            for non_reg_pay in non_reg_payments:
                if self.is_pay_in_regular_date(non_reg_pay):

                    if current_principal <= non_reg_pay.payment: # хотим заплатить больше, чем должны. скорректируем
                        non_reg_pay.payment = current_principal
                    current_principal -= non_reg_pay.payment

                    non_reg_pay.current_principal = current_principal
                    non_reg_pay.principal_payment = non_reg_pay.payment
                    non_reg_pay.interest_payment = 0
                    result.payments[non_reg_pay.date].append(non_reg_pay)

                    if current_principal == 0:
                        break

                    # так как долг поменялся, пересчитаем аннуитентный платёж
                    left_months = self.total_months - m + 1 # сколько месяцев осталось платить
                    monthly_payment = self.calc_monthly_payment(current_principal, left_months)
                else:
                    # считаем набежавшие проценты от последнего платежа до даты досрочки
                    p = self.calc_interest_payment(non_reg_pay.date, cur_period, current_principal)
                    result.year_interest_payments[non_reg_pay.date.year] += p

                    if non_reg_pay.payment < p:
                        pass # не хватает для выплаты процентов
                    else:
                        # обновим остаток долга с учётом процентов

                        if current_principal <= non_reg_pay.payment - p: # хотим заплатить больше, чем должны. скорректируем
                            non_reg_pay.payment = current_principal + p
                        current_principal -= non_reg_pay.payment - p

                        # обновим инфу о платеже и добавим её в список платежей
                        non_reg_pay.interest_payment = p
                        non_reg_pay.current_principal = current_principal
                        non_reg_pay.principal_payment = non_reg_pay.payment - p

                        principal_paid = True # в этом периоде больше за основной долг не платим
                        result.payments[non_reg_pay.date].append(non_reg_pay)

                        if current_principal == 0:
                            break

                        # так как долг поменялся, пересчитаем аннуитентный платёж
                        left_months = self.total_months - m # сколько месяцев осталось платить
                        monthly_payment = self.calc_monthly_payment(current_principal, left_months)

                        cur_period = non_reg_pay.date

            if current_principal == 0:
                break

            # регулярный платёж
            interest_payment = self.calc_interest_payment(end_period, cur_period, current_principal)
            result.year_interest_payments[end_period.year] += interest_payment

            if m == self.total_months and monthly_payment > current_principal: # last payment
                principal_payment = current_principal
                monthly_payment = interest_payment + principal_payment
            elif not principal_paid:
                principal_payment = monthly_payment - interest_payment
            else:
                principal_payment = 0

            current_principal -= principal_payment

            payment = RegularPayment(end_period, interest_payment + principal_payment, interest_payment, principal_payment, current_principal)
            principal_paid = True
            result.payments[payment.date].append(payment)

            if current_principal == 0:
                break

            cur_period = end_period

        return result

def make_non_reg_payment(year, month, day, sum):
    return NonRegularPayment(datetime.date(year, month, day), Decimal(round(sum, 2)), 0, 0, 0)

def add_non_reg_payment(payments, reg_pay_day, year, month, day, sum):
    payments[(year, month if day < reg_pay_day else month + 1)].append(make_non_reg_payment(year, month, day, sum))

def calc_interest_payments(year, periods):
    s = 0
    for p in periods:
        if p.end.year == year:
            for pay in p.payments:
                s += pay.interest_payment

    return s

def diff_dates(dt1, dt2):
    diff = relativedelta(dt1, dt2)
    return diff.years * 12 + diff.months

# каждый месяц добавляем в досрочку разницу между текущим ежемесячным платежём и первоначальным
def calc():

    start_date = datetime.date(2012, 4, 18)
    m = Mortgage(12.4, 12*20, start_date)
    initial_principal = 4900000
    total_months = 12*20
    initial_monthly_payment = m.calc_monthly_payment(initial_principal, total_months)
    next_pay_date = datetime.date(2016, 3, 18)
    non_reg_payments = defaultdict(list)
    add_non_reg_payment(non_reg_payments, start_date.day, 2016, 1, 15, 800000)
    add_non_reg_payment(non_reg_payments, start_date.day, 2016, 2, 18, 10000)

    result = None
    while True:
        result = m.calc(initial_principal, non_reg_payments)
        date, delta, principal = sim(m, initial_monthly_payment, result.payments, next_pay_date)
        if date is None:
            break

        non_reg_payments[(date.year, date.month)].append(NonRegularPayment(next_pay_date, delta, 0, 0, 0))

        # учёт налогового вычета
        if next_pay_date == datetime.date(next_pay_date.year, 5, 18):
            v = round(result.year_interest_payments[next_pay_date.year - 1] * Decimal("0.13"), 2)
            v = math.ceil(v / 1000) * 1000
            non_reg_payments[(date.year, date.month)].append(NonRegularPayment(next_pay_date, v, 0, 0, 0))

        next_pay_date = date

    return result

def sim(mortgage, initial_monthly_payment, payments, date):

    pays = payments.get(date, [])
    for pay in pays:
        if isinstance(pay, RegularPayment):
            delta = initial_monthly_payment - pay.payment
            delta = math.ceil(delta / 1000) * 1000 # округлим до тысячи в большую сторону
            if delta > pay.current_principal: # осталось заплатить меньше, чем получившаяся разница
                delta = pay.current_principal
            next_pay_date = pay.date + relativedelta(months=1)
            return next_pay_date, delta, pay.current_principal
    return None, None, None

def print_payments(payments):
    s1 = 0 # сколько уплатим банку по процентам
    for k, v in sorted(payments.items()):
        for p in v:
            s1 += p.interest_payment
            print(p)
    print("Total interest payment: %s - %% = %s" % (s1, round(s1*Decimal("0.13"), 2)))

if __name__ == '__main__':

    # 2024-09-18 - best result
    # 4587670.65 - total interests payment
    r = calc()
    print_payments(r.payments)
    print("Interest payments by years:")
    for k, v in sorted(r.year_interest_payments.items()):
        print("%s - %s - %% = %s" % (k, v, round(v * Decimal("0.13"), 2)))<commit_msg>fixed add_non_reg_payment function<commit_after># -*- coding: utf-8 -*-

import datetime
from dateutil.relativedelta import relativedelta
import calendar
from itertools import islice
import math
from decimal import Decimal
from collections import defaultdict

class Payment(object):
    def __init__(self, date, payment, interest_payment, principal_payment, current_principal):
        self.date = date
        self.payment = payment
        self.interest_payment = interest_payment
        self.principal_payment = principal_payment
        self.current_principal = current_principal

class RegularPayment(Payment):
    def __init__(self, date, payment, interest_payment, principal_payment, current_principal):
        return super().__init__(date, payment, interest_payment, principal_payment, current_principal)

    def __str__(self):
        return '%s - %.2f - %.2f - %.2f - %.2f - regular' % (str(self.date), self.payment, self.interest_payment, self.principal_payment, self.current_principal)

class NonRegularPayment(Payment):
    def __init__(self, date, payment, interest_payment, principal_payment, current_principal):
        super().__init__(date, payment, interest_payment, principal_payment, current_principal)

    def __str__(self):
        return '%s - %.2f - %.2f - %.2f - %.2f - non regular' % (str(self.date), self.payment, self.interest_payment, self.principal_payment, self.current_principal)

class Period(object):
    def __init__(self, start, end):
        self.start = start
        self.end = end
        self.payments = []
        self.principal_paid = False

    def add_payment(self, payment):
        self.payments.append(payment)

class Result(object):
    def __init__(self):
        self.payments = defaultdict(list)
        self.year_interest_payments = defaultdict(Decimal)

class Mortgage(object):
    def __init__(self, interest, months, start_date):
        self.interest = interest
        self.interest_rate = Decimal(interest / 100)
        self.monthly_interest = Decimal(self.interest_rate / 12)

        # срок кредита
        self.total_months = months

        # дата получения кредита
        self.start_date = start_date

        self.cached_calendar = [self.start_date + relativedelta(months=m) for m in range(0, self.total_months + 1)]

        self.leap_years_cache = [366 if calendar.isleap(self.start_date.year + y) else 365 for y in range(0, int(self.total_months / 12))]

    def calc_monthly_payment(self, principal, months):
        return Decimal(round(principal * (self.monthly_interest / (1 - (1 + self.monthly_interest) ** (- months))), 2))

    def calc_interest_payment(self, pay_date, prev_pay_date, principal):
        # первый месяц в году, т.е. январь
        # в этом случае надо немного по-другому считать
        days_in_year = self.leap_years_cache[self.start_date.year - pay_date.year]
        if pay_date.month == 1 and prev_pay_date.month == 12:
            december31 = datetime.date(prev_pay_date.year, 12, 31)

            days1 = december31 - prev_pay_date
            days1 = days1.days
            p1 = Decimal(days1 / self.leap_years_cache[self.start_date.year - prev_pay_date.year])
            interest_payment1 = Decimal(round(principal * self.interest_rate * p1, 2))

            days2 = pay_date - december31
            days2 = days2.days
            p2 = Decimal(days2 / days_in_year)
            interest_payment2 = Decimal(round(principal * self.interest_rate * p2, 2))

            return interest_payment1 + interest_payment2
        else:
            d = pay_date - prev_pay_date
            days = d.days

            return Decimal(round(principal * self.interest_rate * days / days_in_year, 2))

    def is_pay_in_regular_date(self, pay):
        return pay.date.day == self.start_date.day

    def calc(self, principal, predefined_payments = {}, initial_month = None):
        current_principal = principal
        monthly_payment = self.calc_monthly_payment(current_principal, self.total_months)

        result = Result()
        cur_period = self.cached_calendar[0 if initial_month == None else initial_month]

        for m in range(1 if initial_month == None else initial_month + 1, self.total_months + 1):
            end_period = self.cached_calendar[m]
            principal_paid = False

            non_reg_payments = predefined_payments.get((end_period.year, end_period.month), [])
            for non_reg_pay in non_reg_payments:
                if self.is_pay_in_regular_date(non_reg_pay):

                    if current_principal <= non_reg_pay.payment: # хотим заплатить больше, чем должны. скорректируем
                        non_reg_pay.payment = current_principal
                    current_principal -= non_reg_pay.payment

                    non_reg_pay.current_principal = current_principal
                    non_reg_pay.principal_payment = non_reg_pay.payment
                    non_reg_pay.interest_payment = 0
                    result.payments[non_reg_pay.date].append(non_reg_pay)

                    if current_principal == 0:
                        break

                    # так как долг поменялся, пересчитаем аннуитентный платёж
                    left_months = self.total_months - m + 1 # сколько месяцев осталось платить
                    monthly_payment = self.calc_monthly_payment(current_principal, left_months)
                else:
                    # считаем набежавшие проценты от последнего платежа до даты досрочки
                    p = self.calc_interest_payment(non_reg_pay.date, cur_period, current_principal)
                    result.year_interest_payments[non_reg_pay.date.year] += p

                    if non_reg_pay.payment < p:
                        pass # не хватает для выплаты процентов
                    else:
                        # обновим остаток долга с учётом процентов

                        if current_principal <= non_reg_pay.payment - p: # хотим заплатить больше, чем должны. скорректируем
                            non_reg_pay.payment = current_principal + p
                        current_principal -= non_reg_pay.payment - p

                        # обновим инфу о платеже и добавим её в список платежей
                        non_reg_pay.interest_payment = p
                        non_reg_pay.current_principal = current_principal
                        non_reg_pay.principal_payment = non_reg_pay.payment - p

                        principal_paid = True # в этом периоде больше за основной долг не платим
                        result.payments[non_reg_pay.date].append(non_reg_pay)

                        if current_principal == 0:
                            break

                        # так как долг поменялся, пересчитаем аннуитентный платёж
                        left_months = self.total_months - m # сколько месяцев осталось платить
                        monthly_payment = self.calc_monthly_payment(current_principal, left_months)

                        cur_period = non_reg_pay.date

            if current_principal == 0:
                break

            # регулярный платёж
            interest_payment = self.calc_interest_payment(end_period, cur_period, current_principal)
            result.year_interest_payments[end_period.year] += interest_payment

            if m == self.total_months and monthly_payment > current_principal: # last payment
                principal_payment = current_principal
                monthly_payment = interest_payment + principal_payment
            elif not principal_paid:
                principal_payment = monthly_payment - interest_payment
            else:
                principal_payment = 0

            current_principal -= principal_payment

            payment = RegularPayment(end_period, interest_payment + principal_payment, interest_payment, principal_payment, current_principal)
            principal_paid = True
            result.payments[payment.date].append(payment)

            if current_principal == 0:
                break

            cur_period = end_period

        return result

def make_non_reg_payment(year, month, day, sum):
    return NonRegularPayment(datetime.date(year, month, day), Decimal(round(sum, 2)), 0, 0, 0)

def add_non_reg_payment(payments, reg_pay_day, year, month, day, sum):
    if day < reg_pay_day:
        t = (year, month)
    else:
        d = datetime.date(year, month, 1)
        d += relativedelta(months=1)
        t = (d.year, d.month)
    payments[t].append(make_non_reg_payment(year, month, day, sum))

def calc_interest_payments(year, periods):
    s = 0
    for p in periods:
        if p.end.year == year:
            for pay in p.payments:
                s += pay.interest_payment

    return s

def diff_dates(dt1, dt2):
    diff = relativedelta(dt1, dt2)
    return diff.years * 12 + diff.months

# каждый месяц добавляем в досрочку разницу между текущим ежемесячным платежём и первоначальным
def calc():

    start_date = datetime.date(2012, 4, 18)
    m = Mortgage(12.4, 12*20, start_date)
    initial_principal = 4900000
    total_months = 12*20
    initial_monthly_payment = m.calc_monthly_payment(initial_principal, total_months)
    next_pay_date = datetime.date(2016, 3, 18)
    non_reg_payments = defaultdict(list)
    add_non_reg_payment(non_reg_payments, start_date.day, 2016, 1, 15, 800000)
    add_non_reg_payment(non_reg_payments, start_date.day, 2016, 2, 18, 10000)

    result = None
    while True:
        result = m.calc(initial_principal, non_reg_payments)
        date, delta, principal = sim(m, initial_monthly_payment, result.payments, next_pay_date)
        if date is None:
            break

        non_reg_payments[(date.year, date.month)].append(NonRegularPayment(next_pay_date, delta, 0, 0, 0))

        # учёт налогового вычета
        if next_pay_date == datetime.date(next_pay_date.year, 5, 18):
            v = round(result.year_interest_payments[next_pay_date.year - 1] * Decimal("0.13"), 2)
            v = math.ceil(v / 1000) * 1000
            non_reg_payments[(date.year, date.month)].append(NonRegularPayment(next_pay_date, v, 0, 0, 0))

        next_pay_date = date

    return result

def sim(mortgage, initial_monthly_payment, payments, date):

    pays = payments.get(date, [])
    for pay in pays:
        if isinstance(pay, RegularPayment):
            delta = initial_monthly_payment - pay.payment
            delta = math.ceil(delta / 1000) * 1000 # округлим до тысячи в большую сторону
            if delta > pay.current_principal: # осталось заплатить меньше, чем получившаяся разница
                delta = pay.current_principal
            next_pay_date = pay.date + relativedelta(months=1)
            return next_pay_date, delta, pay.current_principal
    return None, None, None

def print_payments(payments):
    s1 = 0 # сколько уплатим банку по процентам
    for k, v in sorted(payments.items()):
        for p in v:
            s1 += p.interest_payment
            print(p)
    print("Total interest payment: %s - %% = %s" % (s1, round(s1*Decimal("0.13"), 2)))

if __name__ == '__main__':

    # 2024-09-18 - best result
    # 4587670.65 - total interests payment
    r = calc()
    print_payments(r.payments)
    print("Interest payments by years:")
    for k, v in sorted(r.year_interest_payments.items()):
        print("%s - %s - %% = %s" % (k, v, round(v * Decimal("0.13"), 2)))<|endoftext|>


<commit_before>import json
import urllib.request
import zipfile
from os import listdir, remove
from os.path import isfile, join
import re

commands = json.load(open('../data/commands.json', encoding='utf8'))
envs = json.load(open('../data/environments.json', encoding='utf8'))

class PlaceHolder:
    def __init__(self):
        self.count = 0
        self.usePlaceHolders = True

    def setUsePlaceHolders(self, trueOrFalse):
        self.usePlaceHolders = trueOrFalse

    def sub(self, matchObject):
        self.count += 1
        name = ''
        if self.usePlaceHolders:
            name = ':' + matchObject.group(2)
        return  matchObject.group(1) + '${' + str(self.count) + name + '}' + matchObject.group(3)

def get_cwl_files():
    """ Get the list of cwl files from github """
    urllib.request.urlretrieve('https://github.com/LaTeXing/LaTeX-cwl/archive/master.zip', 'cwl.zip')
    zip_ref = zipfile.ZipFile('cwl.zip', 'r')
    zip_ref.extractall('cwl/')
    zip_ref.close()
    remove('cwl.zip')
    files = []
    for f in listdir('cwl/LaTeX-cwl-master'):
        if isfile(join('cwl/LaTeX-cwl-master', f)) and f[-4:] == '.cwl':
            files.append(f)
        else:
            remove(join('cwl/LaTeX-cwl-master', f))
    return files

def parse_cwl_file(file):
    with open(join('cwl/LaTeX-cwl-master', file), encoding='utf8') as f:
        lines = f.readlines()
    pkgcmds = {}
    pkgenvs = []
    for line in lines:
        line = line.rstrip()
        if not line:
            continue
        if line[0] == '#':
            continue
        index_hash = line.find('#')
        if index_hash >= 0:
            line = line[:index_hash]
        if line[:7] == '\\begin{':
            env = line[line.index('{') + 1:line.index('}')]
            if env in envs:
                continue
            pkgenvs.append(env)
            continue
        if line[:5] == '\\end{':
            continue
        if line[0] != '\\':
            continue

        line = line[1:]  # Remove leading '\'
        curly_index = line.find('{')
        square_index = line.find('[')
        # if square_index == -1 and curly_index > -1:
        #     # If there is no optional argument, do not put the arguments in the snippet command
        #     command = line[:curly_index]
        # else:
        command = line
        name = re.sub(r'(\{|\[)[^\{\[\$]*(\}|\])', r'\1\2', command)
        command_dict = {'command': command}
        snippet = line
        if name in commands:
            continue

        p = PlaceHolder()
        if square_index < curly_index:
            # If all the optional args are before {}, we number the {} first
            snippet = re.sub(r'(\{)([^\{\$]*)(\})', p.sub, snippet)
            snippet = re.sub(r'(\[)([^\[\$]*)(\])', p.sub, snippet)
        else:
            snippet = re.sub(r'(\{|\[)([^\{\[\$]*)(\}|\])', p.sub, snippet)

        command_dict['snippet'] = snippet
        pkgcmds[name] = command_dict
    remove(join('cwl/LaTeX-cwl-master', file))
    return (pkgcmds, pkgenvs)

cwl_files = get_cwl_files()
for cwl_file in cwl_files:
    (pkgCmds, pkgEnvs) = parse_cwl_file(cwl_file)
    if pkgEnvs:
        json.dump(pkgEnvs, open(f'../data/packages/{cwl_file[:-4]}_env.json', 'w', encoding='utf8'), indent=2)
    if pkgCmds != {}:
        json.dump(pkgCmds, open(f'../data/packages/{cwl_file[:-4]}_cmd.json', 'w', encoding='utf8'), indent=2)
    # for cmd in pkgCmds:
    #     print(cmd, ': ', pkgCmds[cmd], sep = '')
<commit_msg>Add package name for every command<commit_after>import json
import urllib.request
import zipfile
from os import listdir, remove
from os.path import isfile, join, basename, splitext
import re

commands = json.load(open('../data/commands.json', encoding='utf8'))
envs = json.load(open('../data/environments.json', encoding='utf8'))

class PlaceHolder:
    def __init__(self):
        self.count = 0
        self.usePlaceHolders = True

    def setUsePlaceHolders(self, trueOrFalse):
        self.usePlaceHolders = trueOrFalse

    def sub(self, matchObject):
        self.count += 1
        name = ''
        if self.usePlaceHolders:
            name = ':' + matchObject.group(2)
        return  matchObject.group(1) + '${' + str(self.count) + name + '}' + matchObject.group(3)

def get_cwl_files():
    """ Get the list of cwl files from github """
    urllib.request.urlretrieve('https://github.com/LaTeXing/LaTeX-cwl/archive/master.zip', 'cwl.zip')
    zip_ref = zipfile.ZipFile('cwl.zip', 'r')
    zip_ref.extractall('cwl/')
    zip_ref.close()
    remove('cwl.zip')
    files = []
    for f in listdir('cwl/LaTeX-cwl-master'):
        if isfile(join('cwl/LaTeX-cwl-master', f)) and f[-4:] == '.cwl':
            files.append(f)
        else:
            remove(join('cwl/LaTeX-cwl-master', f))
    return files

def parse_cwl_file(file):
    with open(join('cwl/LaTeX-cwl-master', file), encoding='utf8') as f:
        lines = f.readlines()
    pkgcmds = {}
    pkgenvs = []
    for line in lines:
        line = line.rstrip()
        if not line:
            continue
        if line[0] == '#':
            continue
        index_hash = line.find('#')
        if index_hash >= 0:
            line = line[:index_hash]
        if line[:7] == '\\begin{':
            env = line[line.index('{') + 1:line.index('}')]
            if env in envs:
                continue
            pkgenvs.append(env)
            continue
        if line[:5] == '\\end{':
            continue
        if line[0] != '\\':
            continue

        line = line[1:]  # Remove leading '\'
        curly_index = line.find('{')
        square_index = line.find('[')
        # if square_index == -1 and curly_index > -1:
        #     # If there is no optional argument, do not put the arguments in the snippet command
        #     command = line[:curly_index]
        # else:
        command = line
        name = re.sub(r'(\{|\[)[^\{\[\$]*(\}|\])', r'\1\2', command)
        package = splitext(basename(file))[0]
        command_dict = {'command': command, 'package': package}
        snippet = line
        if name in commands:
            continue

        p = PlaceHolder()
        if square_index < curly_index:
            # If all the optional args are before {}, we number the {} first
            snippet = re.sub(r'(\{)([^\{\$]*)(\})', p.sub, snippet)
            snippet = re.sub(r'(\[)([^\[\$]*)(\])', p.sub, snippet)
        else:
            snippet = re.sub(r'(\{|\[)([^\{\[\$]*)(\}|\])', p.sub, snippet)

        command_dict['snippet'] = snippet
        pkgcmds[name] = command_dict
    remove(join('cwl/LaTeX-cwl-master', file))
    return (pkgcmds, pkgenvs)

cwl_files = get_cwl_files()
for cwl_file in cwl_files:
    (pkgCmds, pkgEnvs) = parse_cwl_file(cwl_file)
    if pkgEnvs:
        json.dump(pkgEnvs, open(f'../data/packages/{cwl_file[:-4]}_env.json', 'w', encoding='utf8'), indent=2)
    if pkgCmds != {}:
        json.dump(pkgCmds, open(f'../data/packages/{cwl_file[:-4]}_cmd.json', 'w', encoding='utf8'), indent=2)
    # for cmd in pkgCmds:
    #     print(cmd, ': ', pkgCmds[cmd], sep = '')
<|endoftext|>


<commit_before>"""
puzzlesolver.py
:created on: 20160624
__author__ = 'Frederic Dupont'
:License: GPL3
"""

__author__ = 'Fred Dupont'

import sys

from sudoku import puzzleconstants as p_const
from sudoku.puzzle import Puzzle, make_grid_from_string

class PuzzleSolver(object):
    """
    class that solves a sudoku puzzle
    takes the clone of a puzzle and solves it
    """
    def __init__(self, puzzle_clone):
        self._puzzle = puzzle_clone

    def _is_solved(self):
        """
        check if the puzzle is solved
        :return: True if solved, False otherwise
        """
        return self._puzzle.is_solved()

    def _is_valid(self):
        """
        check if the puzzle is valid
        :return: True if valid, False otherwise
        """
        return self._puzzle.is_valid()

    def eliminate_candidates(self):
        """For each square in the grid that has a single assigned value,
        run through the PEERS and eliminate this value from the candidates
        *** This will have to be redone each time the puzzle is cloned ***
        :return: False if an inconsistency is discovered --> makes the puzzle invalid
                 True otherwise
        """
        for square in p_const.SQUARES:
            current_value = self._puzzle.grid[square]
            if current_value not in '.0':
                # this says is there is an inconsistency
                # that means it must return False to signal it
                if self._puzzle.candidates[square] != p_const.VALUE_TO_CANDIDATES[current_value]:
                    return False
                for peer in p_const.PEERS[square]:
                    self._puzzle.candidates[peer] = self._puzzle.candidates[peer].replace(current_value, '.')
        return True

    def propagate(self):
        """if a UNIT has only one possible place for a value,
         assign this value there, and adjust the candidates for this place
        """
        for unit in p_const.UNIT_LISTS:
            for digit in p_const.DIGITS:
                result = []
                for square in unit:
                    if digit in self._puzzle.candidates[square]:
                        result.append(square)
                if len(result) == 1:
                    self._puzzle.grid[result[0]] = digit
                    self._puzzle.candidates[result[0]] = p_const.VALUE_TO_CANDIDATES[digit]

    def fill_singles(self):
        """
        completes a unit that is missing only one value
        by updating the grid and the candidates values
        """
        for unit in p_const.UNIT_LISTS:
            result = []
            res_string = [str(_) for _ in range(1, 10)]
            for square in unit:
                val = self._puzzle.grid[square]
                if val in '.0':
                    result.append(square)
                else:
                    res_string[int(val) - 1] = '.'
            if len(result) == 1:
                self._puzzle.grid[result[0]] = p_const.CANDIDATES_TO_VALUE[''.join(res_string)]
                self._puzzle.candidates[result[0]] = ''.join(res_string)

    def eliminate_propagate_fill(self):
        """
        Propagates the constraints by successively eliminating, propagating and filling the
        solved squares until all constraints have been propagated. (no more change in the
        puzzle states)
        ...or the puzzle is solved
        :return: False if an inconsistency is discovered, True otherwise
        """
        while not self._puzzle.is_solved():
            pre_grid_state, pre_candidates_state = repr(self._puzzle), str(self._puzzle)
            if not self.eliminate_candidates():
                return False
            self.propagate()
            self.fill_singles()
            post_grid_state, post_candidates_state = repr(self._puzzle), str(self._puzzle)
            if pre_grid_state == post_grid_state and pre_candidates_state == post_candidates_state:
                return True
        # print(repr(self._puzzle))
        return True

    def search(self):
        """
        clones the puzzle
        creates a new solver
        assigns the next candidate value to the empty square with the less candidates
        recursively calls solve() on the new puzzle
        :return: a solved puzzle repr()
        """
        if self._is_solved():
            return self
        else:
            try:
                new_solver = PuzzleSolver(self._puzzle.clone())
            except AssertionError:
                return False

            next_square = self._get_next_square()
            candidates = [d for d in self._puzzle.candidates[next_square] if d not in '.0']
            for candidate in candidates:

                print('next_square =', next_square, 'candidate = ', candidate)
                new_solver._puzzle.grid[next_square] = candidate
                print('-----------------------> :', end='')
                print(repr(new_solver._puzzle))

                return new_solver.solve()

    def _get_next_square(self):
        """
        Finds the square to explore next:
            --> not filled
            --> with the lowest number of candidates
        :return: the square to explore next
        """
        lowest_num_candidates = 10
        next_square = None
        for open_square in [square for square in p_const.SQUARES if self._puzzle._grid[square] in '.0']:
            num_candidates = sum(1 for d in self._puzzle.candidates[open_square] if d in '123456789')
            if num_candidates == 2:
                return open_square
            elif num_candidates < lowest_num_candidates:
                next_square = open_square
        return next_square

    def solve(self):
        """
        manages the operations to conduct in order to solve a puzzla
        :print: the repr of a solved puzzle (as far as could go with constraint propagation)
        :return: nothing at the moment
        """

        self.eliminate_propagate_fill()
        print('eliminate_propagate_fill :', end='')
        print(repr(self._puzzle))
        self.search()
        print('search                   :', end='')
        print(repr(self._puzzle))
        return repr(self._puzzle)

def main(argv):
    # require search
    g1 = '4.....8.5.3..........7......2.....6.....8.4......1.......6.3.7.5..2.....1.4......'
    partial_s1 = '4.....8.5.3..........7......2.....6.....8.4...4..1.......6.3.7.5.32.1...1.4......'
    s1 = ''
    solve_puzzle(g1)

    g2 = '1...895..5....7819........72.4..8.7.9.71.54.8.8.7..3.531.4..78.4682....3..985...1'
    partial_s2 = '172.895.454..2781989.5142.7254938176937165428681742395315496782468271953729853641'
    s2 = ''
    print()

def solve_puzzle(string):
    print('original string          :', end='')
    print(string)
    solver = PuzzleSolver(make_grid_from_string(string).clone())
    PuzzleSolver.solve(solver)
    print(solver._puzzle.print_puzzle())

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<commit_msg>Add new test requiring search() in main()<commit_after>"""
puzzlesolver.py
:created on: 20160624
__author__ = 'Frederic Dupont'
:License: GPL3
"""

__author__ = 'Fred Dupont'

import sys

from sudoku import puzzleconstants as p_const
from sudoku.puzzle import Puzzle, make_grid_from_string

class PuzzleSolver(object):
    """
    class that solves a sudoku puzzle
    takes the clone of a puzzle and solves it
    """
    def __init__(self, puzzle_clone):
        self._puzzle = puzzle_clone

    def _is_solved(self):
        """
        check if the puzzle is solved
        :return: True if solved, False otherwise
        """
        return self._puzzle.is_solved()

    def _is_valid(self):
        """
        check if the puzzle is valid
        :return: True if valid, False otherwise
        """
        return self._puzzle.is_valid()

    def eliminate_candidates(self):
        """For each square in the grid that has a single assigned value,
        run through the PEERS and eliminate this value from the candidates
        *** This will have to be redone each time the puzzle is cloned ***
        :return: False if an inconsistency is discovered --> makes the puzzle invalid
                 True otherwise
        """
        for square in p_const.SQUARES:
            current_value = self._puzzle.grid[square]
            if current_value not in '.0':
                # this says is there is an inconsistency
                # that means it must return False to signal it
                if self._puzzle.candidates[square] != p_const.VALUE_TO_CANDIDATES[current_value]:
                    return False
                for peer in p_const.PEERS[square]:
                    self._puzzle.candidates[peer] = self._puzzle.candidates[peer].replace(current_value, '.')
        return True

    def propagate(self):
        """if a UNIT has only one possible place for a value,
         assign this value there, and adjust the candidates for this place
        """
        for unit in p_const.UNIT_LISTS:
            for digit in p_const.DIGITS:
                result = []
                for square in unit:
                    if digit in self._puzzle.candidates[square]:
                        result.append(square)
                if len(result) == 1:
                    self._puzzle.grid[result[0]] = digit
                    self._puzzle.candidates[result[0]] = p_const.VALUE_TO_CANDIDATES[digit]

    def fill_singles(self):
        """
        completes a unit that is missing only one value
        by updating the grid and the candidates values
        """
        for unit in p_const.UNIT_LISTS:
            result = []
            res_string = [str(_) for _ in range(1, 10)]
            for square in unit:
                val = self._puzzle.grid[square]
                if val in '.0':
                    result.append(square)
                else:
                    res_string[int(val) - 1] = '.'
            if len(result) == 1:
                self._puzzle.grid[result[0]] = p_const.CANDIDATES_TO_VALUE[''.join(res_string)]
                self._puzzle.candidates[result[0]] = ''.join(res_string)

    def _get_next_square(self):
        """
        Finds the square to explore next:
            --> not filled
            --> with the lowest number of candidates
        :return: the square to explore next
        """
        lowest_num_candidates = 10
        next_square = None
        for open_square in [square for square in p_const.SQUARES if self._puzzle._grid[square] in '.0']:
            num_candidates = sum(1 for d in self._puzzle.candidates[open_square] if d in '123456789')
            if num_candidates == 2:
                return open_square
            elif num_candidates < lowest_num_candidates:
                next_square = open_square
        return next_square

    def eliminate_propagate_fill(self):
        """
        Propagates the constraints by successively eliminating, propagating and filling the
        solved squares until all constraints have been propagated. (no more change in the
        puzzle states)
        ...or the puzzle is solved
        :return: False if an inconsistency is discovered, True otherwise
        """
        while not self._puzzle.is_solved():
            pre_grid_state, pre_candidates_state = repr(self._puzzle), str(self._puzzle)
            if not self.eliminate_candidates():
                return False
            self.propagate()
            self.fill_singles()
            post_grid_state, post_candidates_state = repr(self._puzzle), str(self._puzzle)
            if pre_grid_state == post_grid_state and pre_candidates_state == post_candidates_state:
                return True
        # print(repr(self._puzzle))
        return True

    def search(self):
        """
        clones the puzzle
        creates a new solver
        assigns the next candidate value to the empty square with the less candidates
        recursively calls solve() on the new puzzle
        :return: a solved puzzle repr()
        """
        if self._is_solved():
            return self
        else:
            try:
                new_solver = PuzzleSolver(self._puzzle.clone())
            except AssertionError:
                return False

            next_square = self._get_next_square()
            candidates = [d for d in self._puzzle.candidates[next_square] if d not in '.0']
            for candidate in candidates:

                print('next_square =', next_square, 'candidate = ', candidate)
                new_solver._puzzle.grid[next_square] = candidate
                print('-----------------------> :', end='')
                print(repr(new_solver._puzzle))

                return new_solver.solve()

    def solve(self):
        """
        manages the operations to conduct in order to solve a puzzla
        :print: the repr of a solved puzzle (as far as could go with constraint propagation)
        :return: nothing at the moment
        """

        self.eliminate_propagate_fill()
        print('eliminate_propagate_fill :', end='')
        print(repr(self._puzzle))
        self.search()
        print('search                   :', end='')
        print(repr(self._puzzle))
        return repr(self._puzzle)

def main(argv):

    # # require search
    # g1 = '4.....8.5.3..........7......2.....6.....8.4......1.......6.3.7.5..2.....1.4......'
    # partial_s1 = '4.....8.5.3..........7......2.....6.....8.4...4..1.......6.3.7.5.32.1...1.4......'
    # s1 = ''
    # solve_puzzle(g1)

    # # require search
    # g2 = '1...895..5....7819........72.4..8.7.9.71.54.8.8.7..3.531.4..78.4682....3..985...1'
    # partial_s2 = '172.895.454..2781989.5142.7254938176937165428681742395315496782468271953729853641'
    # s2 = ''
    # solve_puzzle(g2)
    # print()

    #
    g3 = '58261..9.3..79528117928..6....4389..9..126..8..89571..25..61.79.9..72..3....496.2'
    s3 = ''
    solve_puzzle(g3)
    print()

def solve_puzzle(string):
    print('original string          :', end='')
    print(string)
    solver = PuzzleSolver(make_grid_from_string(string).clone())
    PuzzleSolver.solve(solver)
    print(solver._puzzle.print_puzzle())

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<|endoftext|>