<reponame>zig1000/spacechem-level-generator
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import base64
import collections
import copy
import fractions
import gzip
import io
import json

import elements_data

# Element class is defined in elements_data.py to avoid circular dependencies

'''Classes For more apropros errors'''
class FormulaValidityError(ValueError):
    pass
class MoleculeValidityError(ValueError):
    pass

class Formula(collections.Counter):
    '''Represent a chemical formula, as a Counter of elements.'''
    # Redefine Counter's built-in elements() method to return the list of unique ELements in the
    # formula, and move its original functionality to 'atoms()'.
    def elements(self):
        '''Return a list of unique elements in this formula.'''
        # Make sure not to include any 0-counts that Counter leaves hanging around
        return [e for e in self.keys() if self[e] != 0]

    def elements_collection(self):
        '''Return a list containing each element as many times as its count.'''
        return list(super().elements())

    # Have to override Counter's add method or else adding two Formulas will make a Counter
    def __add__(self, other):
        result = Formula()
        for k in self.keys():
            result[k] += self[k]
        for k in other.keys():
            result[k] += other[k]
        return result

    def __mul__(self, other):
        return Formula({i: other * self[i] for i in self.keys()})
    __rmul__ = __mul__

    def num_atoms(self):
        return sum(self.values())

    def least_common_formula(self):
        '''Return a new formula which is this formula divided by the GCD of its element counts'''
        gcd = 0
        for v in self.values():
            gcd = fractions.gcd(v, gcd)

        new_formula = copy.copy(self)
        if gcd > 1: # Make sure we don't divide by 0 in empty case
            for e in new_formula.elements():
                new_formula[e] = new_formula[e] // gcd
        return new_formula

    def get_json_str(self):
        '''Return a string representing this formula using the Hill System (C, then H, then
        alphabetized), in the game's accepted format. E.g. Glycine would be 'C~02H~05NO~02'.
        '''
        result = ''
        # Sort Carbon and Hydrogen to the front and alphabetize the rest
        elements = sorted(self.elements(),
                          key=lambda e: '0' if e.symbol == 'C'
                                        else '1' if e.symbol == 'H'
                                        else e.symbol)
        for e in elements:
            result += e.symbol
            if self[e] != 1:
                result += '~' + str(self[e]).rjust(2, '0')
        return result
    __str__ = get_json_str # For debugging convenience

    def is_valid(self, large_output=False):
        '''Check if it's possible to form a molecule with this formula within an input/output zone.
        Empty formulas are considered invalid. Default 4x4 zone, optionally large (8x4) output zone.
        '''
        # Verify size constraints
        if not 1 <= self.num_atoms() <= 16 + 16*large_output:
            return False

        # We'll calculate validity based on whether there are enough bonds in the atom list to form
        # a minimally-connected molecule. To check this, consider a simple linearly connected
        # molecule with 2 endpoints. An "endpoint" element has max bonds 1, while elements with 3 or
        # 4+ max bonds each allow for 1 or 2 additional endpoints in the molecule, respectively.
        # Elements of max bonds 2 do not affect the count and thus do not affect the validity of
        # the formula, apart from the 16 or 32-atom limit of the zone.
        #
        # Though not formally proven, it appears that the maximum number of endpoints is in each
        # case a constant equal to half the number of cells in the zone. By trial and error with the
        # 4x4 zone it appears that this maximum can be reached with any composition of 3 vs 4 bond
        # atoms, and a simplifying assumption is made that this holds for the 8x4 case.
        # The cases for which an incorrect return value by this method could cause an exception
        # would in any case be prohibitively rare and of little concern, and imo worth the tradeoff
        # for O(k) runtime (where k is the # of unique elements in the formula).

        # Due to the limited sizes of the zones, each max bound count element can only contribute
        # extra connections so many times before it reaches its limit. For example, in a 4x4 zone,
        # only two atoms with max bond count 4 will provide two extra endpoints each. Subsequent
        # max bond count 4 atoms will only allow for 1 extra endpoint each due to space
        # constraints. E.g. C3H8 is impossible to construct in a 4x4 zone.
        #
        if large_output:
            extra_endpoints_dict = {3:14, 4:6}
        else:
            extra_endpoints_dict = {3:6, 4:2}

        allowed_endpoint_count = 2 # H-H base case
        for element in self.elements():
            if element.max_bonds == 0:
                # A noble gas is only valid as the only atom in a molecule
                return self.num_atoms() == 1
            elif element.max_bonds == 1:
                allowed_endpoint_count -= self[element]
            # Count one extra allowed endpoint per atom with bond count 3 or more (up to limit)
            elif element.max_bonds >= 3:
                # As long as our formula has no negatives this should be fine
                extra_endpoints = min(self[element], extra_endpoints_dict[3])
                allowed_endpoint_count += extra_endpoints
                extra_endpoints_dict[3] -= extra_endpoints

                # Count an additional extra endpoint per atom with bond count 4 or more (up to limit)
                if element.max_bonds >= 4:
                    extra_endpoints = min(self[element],
                                          extra_endpoints_dict[3],
                                          extra_endpoints_dict[4])
                    allowed_endpoint_count += extra_endpoints
                    extra_endpoints_dict[3] -= extra_endpoints
                    extra_endpoints_dict[4] -= extra_endpoints

        return allowed_endpoint_count >= 0

    def fission_sources(self):
        '''Return a dict of atomic masses and their counts for all elements that could have fission
        performed on them to obtain part of this formula.
        The heavy lifting of this method is tucked away at the bottom of this file in
        splittable_sources since it's a monster of a function.
        '''
        output_masses = collections.Counter({e.atomic_num: count for e, count in self.items()
                                             if count > 0})
        return splittable_sources(output_masses)

    def remove_fissile_element(self, element, count):
        '''Used while converting an output formula to an input formula via inverse fission.
        Given a fissile element and count, remove the target count of the element from
        this formula, drilling down into its 'fission tree' as needed.
        Raise FormulaValidityError if the element / its fission tree doesn't add up to the count.
        '''
        # Remove as much as we can of this element without fission
        direct_removals = min(count, self[element])
        count -= direct_removals
        self[element] -= direct_removals
        # Keep this object clean
        if self[element] == 0:
            del self[element]

        if count != 0:
            # If we hit the bottom of the fission tree and aren't done, raise an exception
            if element.atomic_num == 1:
                raise FormulaValidityError(f"Couldn't remove {count} of {element} from formula")

            try:
                if element.atomic_num % 2 == 0:
                    child = elements_data.elements_dict[element.atomic_num // 2]
                    self.remove_fissile_element(child, 2*count)
                else:
                    child_A, child_B = (elements_data.elements_dict[element.atomic_num // 2 + 1],
                                        elements_data.elements_dict[element.atomic_num // 2])
                    self.remove_fissile_element(child_A, count)
                    self.remove_fissile_element(child_B, count)
            except FormulaValidityError:
                raise FormulaValidityError(f"Couldn't remove {count} of {element} from formula")

# Enum-esque directional vars for convenience
DIRECTIONS = UP, RIGHT, DOWN, LEFT = (0, 1, 2, 3) # Python, baby

def opposite_dir(dir):
    '''Given an Int representing a direction return its opposite direction.'''
    return (dir + 2) % 4

class GridPos:
    '''Represent a 0-indexed (row, col) position within an input/output zone.
    Indices increase from left to right and top to bottom.
    '''
    num_cols = 4

    def __init__(self, row, col, large_output=False):
        self.row = row
        self.col = col
        self.large_output = large_output
        self.num_rows = 4 + 4*large_output

    def __str__(self):
        return f'({self.row}, {self.col})'
    __repr__ = __str__

    # __eq__ and __hash__ so we can use GridPos as dictionary keys
    def __eq__(self, other):
        return type(self) == type(other) and (self.row, self.col) == (other.row, other.col)

    def __hash__(self):
        return hash((self.row, self.col))

    def is_valid(self):
        '''Check that this position consists of integer positions within the zone's grid.'''
        return isinstance(self.row, int) and isinstance(self.col, int) \
               and (0 <= self.row < self.num_rows) and (0 <= self.col < self.num_cols)

    def neighbor(self, dir):
        '''Return the neighbor GridPos in the indicated direction, or None if out-of-bounds.'''
        if dir == UP:
            r, c = self.row - 1, self.col
        elif dir == RIGHT:
            r, c = self.row, self.col + 1
        elif dir == DOWN:
            r, c = self.row + 1, self.col
        elif dir == LEFT:
            r, c = self.row, self.col - 1
        else:
            raise ValueError(f"Invalid direction: {dir}")

        if 0 <= r < self.num_rows and 0 <= c < self.num_cols:
            return GridPos(r, c, self.large_output)
        return None

    def neighbors(self):
        '''Return all orthogonally adjacent positions within the zone's grid.'''
        return [p for p in (self.neighbor(dir) for dir in DIRECTIONS) if p is not None]

    def dirs_and_neighbors(self):
        '''Return a list of (dir, pos) pairs for each neighboring position within the grid.'''
        return [(d, p) for d, p in ((_d, self.neighbor(_d)) for _d in DIRECTIONS) if p is not None]

class Atom:
    '''Represent an Atom, including its element, grid position, and attached bonds.
    '''
    def __init__(self, element, pos):
        self.bonds = [0, 0, 0, 0] # up, right, down, left
        self.set_element(element)
        self.set_pos(pos)

    def __str__(self):
        return self.symbol.rjust(2) # Pad element symbol to two chars

    def __repr__(self):
        return f'Atom({self.symbol}, {self.pos}, {self.bonds})'

    def __eq__(self, other):
        return (type(self) == type(other)
                and self.element == other.element
                and self.pos == other.pos
                and self.bonds == other.bonds)

    def get_json_str(self):
        '''Return a string representing this atom in the level json's format.'''
        return f'{self.col}{self.row}{self.atomic_num}{self.bonds[RIGHT]}{self.bonds[DOWN]}'

    def remaining_bonds(self):
        '''Return the # of remaining bonds this atom is allowed.'''
        return self.max_bonds - sum(self.bonds)

    def set_pos(self, pos):
        '''Change this atom's position in the grid.'''
        self.pos = pos
        self.row = self.pos.row
        self.col = self.pos.col

    def set_element(self, element):
        if sum(self.bonds) > element.max_bonds:
            raise ValueError(f"Too many bonds to change atom {self} to element {element}")

        self.element = element

        # Exposing some sub-attributes for convenience
        self.atomic_num = element.atomic_num
        self.symbol = element.symbol
        self.max_bonds = element.max_bonds

class Molecule:
    '''Represents an input/output zone and the molecule constructed therein.
    '''
    def __init__(self, large_output=False):
        self.name = 'Randite'
        self.large_output = large_output
        self.num_rows = 4 + 4*large_output
        self.num_cols = 4
        self.formula = Formula()
        # TODO: Could potentially merge grid and used_posns into a single GridPos:Atom dict.
        #       The main convenience of the grid at this point is that it'll automatically yell at us
        #       if we start asking for atoms from a GridPos that's out-of-bounds.
        self.grid = [[None, None, None, None] for _ in range(self.num_rows)]
        self.used_posns = set() # Tracked so that we can easily iterate over the atoms in the molecule

        # To optimize the performance of available_positions(), we'll roughly track the # of open
        # bonds available on this molecule.
        # An atom with no open adjacencies in the grid contributes 0 to this count.
        self.open_bonds = 0

    def __getitem__(self, pos):
        '''Return the atom at the specified grid position or None.'''
        return self.grid[pos.row][pos.col]

    def __setitem__(self, pos, item):
        '''Set the specified grid position (item should be None or an Atom).'''
        self.grid[pos.row][pos.col] = item
        if item is None:
            self.used_posns.remove(pos)
        else:
            self.used_posns.add(pos)

    def __iter__(self):
        '''Iterate over each atom in this molecule. Order of iteration is not defined.'''
        return (self[p] for p in self.used_posns)

    def __len__(self):
        '''Return the # of atoms in this molecule.'''
        return len(self.used_posns)

    def __str__(self):
        '''Pretty-print this molecule.'''
        result = ' _________________ \n' # Border of the input/output zone
        for r in range(self.num_rows):
            result += '|'
            for c in range(self.num_cols):
                atom = self.grid[r][c]
                # Represent any atoms here
                if atom is None:
                    result += 2*' '
                else:
                    result += str(atom).rjust(2)
                # Represent any bonds to the right of the atom
                left_atom = atom
                right_atom = self.grid[r][c + 1] if c + 1 < self.num_cols else None

                bond_str = ' '
                if left_atom is not None and right_atom is not None \
                   and left_atom.bonds[RIGHT] != right_atom.bonds[LEFT]:
                    bond_str = '?'
                elif left_atom is not None and left_atom.bonds[RIGHT] != 0:
                    bond_str = str(left_atom.bonds[RIGHT])
                elif right_atom is not None and right_atom.bonds[LEFT] != 0:
                    bond_str = str(right_atom.bonds[LEFT])
                if c < self.num_cols - 1:
                    result += ' ' + bond_str + ' '
            result += '|\n'
            # Add a row of vertical bonds
            if r < self.num_rows - 1:
                result += '|'
                for c in range(self.num_cols):
                    top_atom = self.grid[r][c]
                    if r + 1 < self.num_rows:
                        bottom_atom = self.grid[r + 1][c]
                    else:
                        bottom_atom = None
                    bond_str = '  '
                    if top_atom is not None and bottom_atom is not None \
                       and top_atom.bonds[DOWN] != bottom_atom.bonds[UP]:
                        bond_str = '??'
                    elif top_atom is not None and top_atom.bonds[DOWN] != 0:
                        bond_str = ' ' + str(top_atom.bonds[DOWN])
                    elif bottom_atom is not None and bottom_atom.bonds[UP] != 0:
                        bond_str = ' ' + str(bottom_atom.bonds[UP])
                    result += bond_str
                    if c < self.num_cols - 1:
                        result += 3*' '
                result += '|\n'
        result += '|_________________|\n'
        return result

    __repr__ = __str__

    def get_json_str(self):
        '''Return a string representing this molecule in the level json's format.'''
        result = f'{self.name};{self.formula.get_json_str()}'
        for atom in self:
            result += ';' + atom.get_json_str()
        return result

    def update_formula(self):
        '''To be called after mutating any atom elements. Update the formula of this molecule.'''
        self.formula = Formula()
        for atom in self:
            self.formula[atom.element] += 1

    def update_open_bonds(self):
        '''Update the count of open bonds. Since we only care about updating it well
        enough to know when it's 0, we'll ignore the triple bond limit, and count any open side of
        an atom as adding the remainder of its max bond count to the open bonds.
        '''
        self.open_bonds = 0
        for atom in self:
            if any(self[pos] is None for pos in atom.pos.neighbors()):
                self.open_bonds += atom.remaining_bonds() # Not exact but we don't need it to be

    def open_positions(self):
        '''Return a list of valid grid positions where an atom could be added to this molecule.'''
        # For an empty molecule, all positions are open
        if len(self) == 0:
            return [GridPos(r, c, large_output=self.large_output)
                    for r in range(self.num_rows) for c in range(self.num_cols)]
        # If there are no remaining bonds, we can skip the overhead of walking through the atoms
        elif self.open_bonds == 0:
            return []

        checked_posns = set() # For O(1) checks on whether a position has already been added
        for atom in self:
            if atom.remaining_bonds() > 0:
                for pos in atom.pos.neighbors():
                   if self[pos] is None and pos not in checked_posns:
                      checked_posns.add(pos)
        return list(checked_posns)

    def add_atom(self, new_atom):
        '''Adds the given Atom to this molecule. The Atom's position must be open in this molecule.
        Also adds any bonds specified by the incoming atom to its neighboring atoms.
        For convenience of more complex operations, it is allowable to add an atom with unfulfilled
        bonds or which is not connected to the rest of the molecule.
        '''
        if self[new_atom.pos] is not None:
            raise Exception(f"Conflict with existing atom; cannot add {repr(new_atom)} to \n{self}")

        # Add the atom into our grid / formula. Then add its bonds while re-calculating self.open_bonds
        self[new_atom.pos] = new_atom
        self.used_posns.add(new_atom.pos)
        self.formula[new_atom.element] += 1

        # Quick helper to check if an atom within this molecule's grid has at least 1 open side
        def has_open_side(atom):
            return any(self[pos] is None for pos in atom.pos.neighbors())

        # Partial update of the number of open bonds this molecule has
        if has_open_side(new_atom):
            self.open_bonds += new_atom.remaining_bonds()

        # Add bonds to neighbours matching the bonds indicated on this atom
        for dir, pos in new_atom.pos.dirs_and_neighbors():
            adj_atom = self[pos]
            if adj_atom is not None:
                adj_atom.bonds[opposite_dir(dir)] = new_atom.bonds[dir]
                # Subtract the bond we just added from the molecule's 'open bonds'
                self.open_bonds -= new_atom.bonds[dir]

                # If we closed off the neighbor's last open face, we've additionally removed
                # however many bonds it now has left from the molecule's 'open' bonds
                if not has_open_side(adj_atom):
                    self.open_bonds -= adj_atom.remaining_bonds()

    def remove_atom(self, atom):
        '''Remove the specified atom from this molecule. Must exactly match.'''
        if self[atom.pos] != atom:
            raise ValueError(f"Specified atom {repr(new_atom)} does not match an atom in:\n{self}"
                             + "\nCannot be removed.")

        self[atom.pos] = None
        self.formula[atom.element] -= 1

        # Remove any now-trailing bonds on neighbors
        for dir, pos in atom.pos.dirs_and_neighbors():
            adj_atom = self[pos]
            if adj_atom is not None:
                adj_atom.bonds[opposite_dir(dir)] = 0
        self.update_open_bonds()

    def is_connected(self):
        '''For the purposes of more advanced construction algorithms we allow adding atoms in
        unconnected cells. This checks if the molecule is currently 'connected' and thus valid.
        We'll count empty molecules as unconnected.
        '''
        if len(self) == 0:
            return False

        # Do a DFS starting from one atom and following the bonds of the molecule. If we don't
        # find every atom, it's not connected
        sample_pos = next(iter(self.used_posns))
        stack = [self[sample_pos]]
        # We don't have to actually 'visit' every atom, seeing them as neighbors is sufficient
        seen = {sample_pos} # Track the grid positions of seen connected atoms
        while stack:
            if len(seen) == len(self):
                return True

            atom = stack.pop()
            # Check for connected neighbors. When we see an unseen connected atom, add it to the
            # stack
            for dir, adj_pos in atom.pos.dirs_and_neighbors():
                if atom.bonds[dir] != 0 and adj_pos not in seen:
                    seen.add(adj_pos)
                    adj_atom = self[adj_pos]
                    stack.append(adj_atom)
        return False

    def shift(self, rows=0, cols=0):
        '''Shift the current contents of this molecule downward/rightward by the specified number
        of rows/columns. Negative numbers shift upward/leftward.

        Raise an exception if this would place atoms out-of-bounds.
        '''
        # Make sure this is a legal shift
        for atom in self:
            if (atom.row + rows < 0 or atom.row + rows > self.num_rows) \
               or (atom.col + cols < 0 or atom.col + cols > self.num_cols):
               raise Exception(f'Cannot shift molecule\n{self}\nby {rows} rows and {cols} cols')

        # Wipe the grid clean and re-add the atoms in their new positions
        atoms = list(self)
        self.grid = [[None, None, None, None] for r in range(self.num_rows)]
        self.used_posns = set()
        for atom in atoms:
            atom.set_pos(GridPos(atom.row + rows, atom.col + cols, large_output=self.large_output))
            self[atom.pos] = atom

        # Recount open bonds once we're done since some atoms may no longer have open sides
        self.update_open_bonds()

    def add_molecule(self, other):
        '''Add the specified molecule to this molecule. Must not have any atoms in conflicting
        positions.
        '''
        # Check for conflicts
        if any(self[atom.pos] is not None for atom in other):
            raise Exception(f'Cannot add molecule\n{other} to molecule\n{self}; conflicting atoms')

        # Add the new atoms
        for atom in other:
            self.add_atom(atom)

class Level:
    '''Parent class for Research and Production levels.'''
    def __init__(self):
        self.dict = {}

    def __getitem__(self, item):
        return self.dict[item]

    def __setitem__(self, item, val):
        self.dict[item] = val

    def __str__(self):
        return json.dumps(self.dict)

    def get_code(self):
        '''Get the mission code - gzip then b64 the level json.'''
        out = io.BytesIO()
        with gzip.GzipFile(fileobj=out, mode="w") as f:
            f.write(json.dumps(self.dict).encode('utf-8'))
        return base64.b64encode(out.getvalue()).decode()

class ResearchLevel(Level):
    def __init__(self):
        Level.__init__(self)
        self['input-zones'] = {}
        self['output-zones'] = {}

        self['has-large-output'] = False

        # Features of the level
        self['bonder-count'] = 0
        self['has-sensor'] = False
        self['has-fuser'] = False
        self['has-splitter'] = False
        self['has-teleporter'] = False

        self['type'] = 'research'
        self['name'] = 'RandomlyGenerated'
        self['author'] = "<NAME>"
        self['difficulty'] = 0

class ProductionLevel(Level):
    def __init__(self):
        Level.__init__(self)
        self['terrain'] = 0
        self['random-input-zones'] = {} # Max 1 random
        self['fixed-input-zones'] = {} # Max 2 fixed
        self['output-zones'] = {} # Max 3 outputs

        self['max-reactors'] = 6 # Default maximum allowed

        self['has-starter'] = False
        self['has-assembly'] = False
        self['has-disassembly'] = False
        self['has-advanced'] = False # Sensor reactor
        self['has-nuclear'] = False
        self['has-superbonder'] = False
        self['has-recycler'] = False

        self['type'] = 'production'
        self['name'] = 'RandomlyGenerated'
        self['author'] = "<NAME>"
        self['difficulty'] = 0

def splittable_sources(given):
    '''Given a Counter of ints, return a dict of ints and their total counts that can be
    non-trivially constructed from the given integers, using only series' of addition of integers
    within 1 of each other, and where no given int is used more times than its count in the Counter.
    In other words, we're using the inverse fission operation to calculate viable input elements
    that could have been split any (non-0) # of times to create part of the given output.
    '''
    # NOTE: We ask for a Counter as input because unlike dicts, they implicitly return 0 if
    #       asked for a key they don't contain, which simplifies our code

    # Tally tracking ints we were given/discover, the max of each we can create at once, and
    # some additional helper values - as needed we'll create dicts that track how many N-1 ints
    # we can create at the same time as any possible count of N's. We'll also track the
    # 'most balanced' such allocation of N vs N-1, which will allow us to properly get counts for
    # odd ints, as well as assisting in the creation of higher dicts.
    tally = {}
    # Dict of ints that were constructable (not just given), and their max counts
    constructed = {}

    # Keep a running sum of what we were given so we don't waste time on clearly impossible sums
    givens_sum = 0

    # To avoid the overhead of a priority queue, use one queue for the given ints,
    # and one queue for ints we obtained by addition (we'll add them in numerical order).
    # Loop on whichever's next value is lower
    given_queue = collections.deque(sorted(given.keys()))
    added_queue = collections.deque()
    while given_queue or added_queue:
        # Pop the element we're iterating on - pop from both queues at once if they match
        if (not added_queue
            or (given_queue and given_queue[0] < added_queue[0])):
            n = given_queue.popleft()
        else:
            if given_queue and given_queue[0] == added_queue[0]:
                given_queue.popleft()
            n = added_queue.popleft()

        # Calculate how many copies of n we can obtain at once
        if n % 2 == 0:
            # If n is even, we only need to update its count, based on
            # the count of n / 2 and however much of n we were given to start
            component_count = tally[n // 2]['count'] if n // 2 in tally else 0
            this_count = component_count // 2 + given[n]
        else:
            # To count odd n, we must first make a dict that pairs the max # of n // 2 that can
            # be created for any given count of n // 2 + 1. We can do this recursively off
            # a previous such dict. When creating this dict we will also store the count
            # of n // 2 + 1 for which there can simultaneously be created a most closely balanced
            # count of n // 2. This can be used directly to count n and also to speed up
            # the recursive creation of dicts.
            # However we can skip this if either of n // 2 or n // 2 + 1 are unavailable.
            # Note that even if both are available they may not be addable so our count could
            # still come out to 0
            upper_child, lower_child = n // 2 + 1, n // 2
            if upper_child in tally and lower_child in tally:
                # In this case, calculate and store upper_child's neighbour dict
                tally[upper_child]['neighbour_counts'] = {}
                if upper_child == 2: # Calc 2->1 dict
                    balanced_upper_count, min_count_in_best_balance = -1, -1
                    for upper_count in range(1, tally[upper_child]['count'] + 1):
                        lower_count = (tally[lower_child]['count']
                                       - 2*max(upper_count - given[upper_child], 0))
                        tally[upper_child]['neighbour_counts'][upper_count] = lower_count

                        # Check how balanced this allocation is
                        worst_count = min(upper_count, lower_count)
                        if worst_count >= min_count_in_best_balance:
                            balanced_upper_count = upper_count
                            min_count_in_best_balance = worst_count
                    tally[upper_child]['balanced_count'] = balanced_upper_count
                elif upper_child == 3: # Calc 3->2 dict
                    balanced_upper_count, min_count_in_best_balance = -1, -1
                    for upper_count in range(1, tally[upper_child]['count'] + 1):
                        nongiven_upper_count = max(upper_count - given[upper_child], 0)
                        # 2s count = (2s constructable given 1s used in 3s) - (2s used in 3s)
                        lower_count = (given[2] + (given[1] - nongiven_upper_count) // 2
                                       - nongiven_upper_count)
                        tally[upper_child]['neighbour_counts'][upper_count] = lower_count

                        # Check how balanced this allocation is
                        worst_count = min(upper_count, lower_count)
                        if worst_count >= min_count_in_best_balance:
                            balanced_upper_count = upper_count
                            min_count_in_best_balance = worst_count
                    # Store the most balanced count for upper_child
                    tally[upper_child]['balanced_count'] = balanced_upper_count
                # If either the upper child or the lower child had no compound components,
                # the upper_child's neighbour_counts dict is just the max count of lower_child,
                # regardless of the count of upper_child
                elif (tally[upper_child]['count'] == given[upper_child]
                      or tally[lower_child]['count'] == given[lower_child]):
                    tally[upper_child]['neighbour_counts'] = {
                            i: tally[lower_child]['count']
                            for i in range(1, tally[upper_child]['count'] + 1) }
                    # Since the lower_child gets the same count no matter what, just maximize
                    # upper_child's count for the 'balanced' allocation
                    tally[upper_child]['balanced_count'] = tally[upper_child]['count']
                # Otherwise, based on our recursion principle, the upper child's upper
                # child must already have its neighbour_counts dict set. Use that to calculate
                # the upper child's neighbour_counts. The algorithm for this depends on which of
                # upper_child/lower_child is even.
                # We also have a couple of base cases to handle when building the neighbour dict
                # dict for 3->2 and 2-> 1, since in those cases lower_child is also a component
                # of upper_child
                elif upper_child % 2 == 0:
                    # If the upper child is even, calculate how much of lower_child's components
                    # are used up by any valid count of upper_child, and thus the max
                    # lower_child count for that count of upper_child.
                    # Call A upper_child // 2 and B the other (lower) component of lower_child
                    A = upper_child // 2
                    balanced_upper_count, min_count_in_best_balance = -1, -1
                    for upper_count in range(1, tally[upper_child]['count'] + 1):
                        if upper_count <= given[upper_child]:
                            lower_count = tally[lower_child]['count']
                        else:
                            A_used_in_upper = 2*(upper_count - given[upper_child])
                            if A_used_in_upper == tally[A]['count']:
                                lower_count = given[lower_child]
                            else:
                                # Search to the right of the original balance point and/or our
                                # new A limit, to find a balance given the unusable As:
                                start_idx = max(tally[A]['balanced_count'], A_used_in_upper + 1)
                                built_lower_count = 0
                                for used_A in range(start_idx, tally[A]['count'] + 1):
                                    worst_count = min(used_A - A_used_in_upper,
                                                      tally[A]['neighbour_counts'][used_A])
                                    built_lower_count = max(built_lower_count, worst_count)
                                lower_count = built_lower_count + given[lower_child]
                        tally[upper_child]['neighbour_counts'][upper_count] = lower_count

                        # Check how balanced this allocation is
                        worst_count = min(upper_count, lower_count)
                        if worst_count >= min_count_in_best_balance:
                            balanced_upper_count = upper_count
                            min_count_in_best_balance = worst_count
                    # Store the most balanced count for upper child
                    tally[upper_child]['balanced_count'] = balanced_upper_count
                else:
                    # If the upper child is odd, call its subchildren A and B, where B =
                    # lower_child // 2. Using A's neighbour_counts dict, calculate how much B and
                    # from that how much lower_child we can make for any valid count of
                    # upper_child
                    A = upper_child // 2 + 1
                    # For each possible count of upper_child, calculate how many copies of
                    # lower_child can be simultaneously constructed from the leftovers
                    # Also track the 'most balanced' count we can assign to upper_child
                    balanced_upper_count, min_count_in_best_balance = -1, -1
                    for upper_count in range(1, tally[upper_child]['count'] + 1):
                        if upper_count <= given[upper_child]:
                            lower_count = tally[lower_child]['count']
                        else:
                            used_A = used_B = upper_count - given[upper_child]
                            available_B = (tally[A]['neighbour_counts'][used_A] - used_A)
                            lower_count = available_B // 2 + given[lower_child]
                        tally[upper_child]['neighbour_counts'][upper_count] = lower_count

                        # Check how balanced this allocation is
                        worst_count = min(upper_count, lower_count)
                        if worst_count >= min_count_in_best_balance:
                            balanced_upper_count = upper_count
                            min_count_in_best_balance = worst_count
                    # Store the most balanced count for upper child
                    tally[upper_child]['balanced_count'] = balanced_upper_count

                # Calculate the count of n based on upper_child's most balanced count
                balanced_upper_count = tally[upper_child]['balanced_count']
                this_count = (min(balanced_upper_count,
                                  tally[upper_child]['neighbour_counts'][balanced_upper_count])
                              + given[n])
            else:
                # If n only occurred as an input and not a compound, set it to the given count
                # The n = 1 case is handled here since 1 can never be compound
                # We don't need to calculate its neighbour dict in this case.
                this_count = given[n]

        # If the count came out to 0, ignore this int
        if this_count == 0:
            continue
        # Update the tally with the discovered count
        tally[n] = {'count': this_count}
        #  Add this int to the output dict if it was possible to construct
        # (not just obtained from the givens)
        if this_count != given[n]:
            constructed[n] = this_count

        # Add any viable sums (restricted to valid atomic masses) obtained from n to the queue
        # As a mini-optimization, we won't add odd numbers to the queue that exceed the sum of
        # the givens up to n
        givens_sum += n*given[n]
        # If n - 1 is in the tally, add 2n - 1 to the queue
        if (n - 1 in tally
            and (2*n - 1 <= 109 or 2*n - 1 in (201, 203))
            and 2*n - 1 <= givens_sum):
            added_queue.append(2*n - 1)
        # If the count for n was at least 2, add 2n to the queue
        if tally[n]['count'] >= 2 and (2*n <= 109 or 2*n in (200, 202)):
            added_queue.append(2*n)

    # Once we've looped over all possible sums, return a dict of the relevant ints and their counts
    return constructed


<filename>src/hcb/tools/analysis/probability_util.py
import math
from typing import Union, Callable, Sequence, Tuple

import numpy as np
from scipy.stats import linregress
from scipy.optimize import leastsq
from scipy.stats._stats_mstats_common import LinregressResult

def log_binomial(*, p: Union[float, np.ndarray], n: int, hits: int) -> Union[float, np.ndarray]:
    r"""Approximates $\ln(P(hits = B(n, p)))$; the natural logarithm of a binomial distribution.

    All computations are done in log space to ensure intermediate values can be represented as
    floating point numbers without underflowing to 0 or overflowing to infinity. This is necessary
    when computing likelihoods over many samples. For example, if 80% of a million samples are hits,
    the maximum likelihood estimate is p=0.8. But even this optimal estimate assigns a prior
    probability of roughly 10^-217322 for seeing *exactly* 80% hits out of a million (whereas the
    smallest representable double is roughly 10^-324).

    This method can be broadcast over multiple hypothesis probabilities.

    Args:
        p: The independent probability of a hit occurring for each sample. This can also be an array
            of probabilities, in which case the function is broadcast over the array.
        n: The number of samples that were taken.
        hits: The number of hits that were observed amongst the samples that were taken.

    Returns:
        $\ln(P(hits = B(n, p)))$
    """
    # Clamp probabilities into the valid [0, 1] range (in case float error put them outside it).
    p_clipped = np.clip(p, 0, 1)

    result = np.zeros(shape=p_clipped.shape, dtype=np.float32)
    misses = n - hits

    # Handle p=0 and p=1 cases separately, to avoid arithmetic warnings.
    if hits:
        result[p_clipped == 0] = -np.inf
    if misses:
        result[p_clipped == 1] = -np.inf

    # Multiply p**hits and (1-p)**misses onto the total, in log space.
    result[p_clipped != 0] += np.log(p_clipped[p_clipped != 0]) * hits
    result[p_clipped != 1] += np.log1p(-p_clipped[p_clipped != 1]) * misses

    # Multiply (n choose hits) onto the total, in log space.
    log_n_choose_hits = log_factorial(n) - log_factorial(misses) - log_factorial(hits)
    result += log_n_choose_hits

    return result

def log_factorial(n: int) -> float:
    r"""Approximates $\ln(n!)$; the natural logarithm of a factorial.

    Uses Stirling's approximation for large n.
    """
    if n < 20:
        return sum(math.log(k) for k in range(1, n + 1))
    return (n + 0.5) * math.log(n) - n + math.log(2 * np.pi) / 2

def binary_search(*, func: Callable[[int], float], min_x: int, max_x: int, target: float) -> int:
    """Performs an approximate granular binary search over a monotonically ascending function."""
    while max_x > min_x + 1:
        med_x = (min_x + max_x) // 2
        out = func(med_x)
        if out < target:
            min_x = med_x
        elif out > target:
            max_x = med_x
        else:
            return med_x
    fmax = func(max_x)
    fmin = func(min_x)
    dmax = 0 if fmax == target else fmax - target
    dmin = 0 if fmin == target else fmin - target
    return max_x if abs(dmax) < abs(dmin) else min_x

def binary_intercept(*, func: Callable[[float], float], start_x: float, step: float, target_y: float, atol: float) -> float:
    """Performs an approximate granular binary search over a monotonically ascending function."""
    start_y = func(start_x)
    if abs(start_y - target_y) <= atol:
        return start_x
    while (func(start_x + step) >= target_y) == (start_y >= target_y):
        step *= 2
        if np.isinf(step) or step == 0:
            raise ValueError("Failed.")
    xs = [start_x, start_x + step]
    min_x = min(xs)
    max_x = max(xs)
    increasing = func(min_x) < func(max_x)

    while True:
        med_x = (min_x + max_x) / 2
        med_y = func(med_x)
        if abs(med_y - target_y) <= atol:
            return med_x
        assert med_x not in [min_x, max_x]
        if (med_y < target_y) == increasing:
            min_x = med_x
        else:
            max_x = med_x

def least_squares_cost(*, xs: np.ndarray, ys: np.ndarray, intercept: float, slope: float) -> float:
    assert len(xs.shape) == 1
    assert xs.shape == ys.shape
    return np.sum((intercept + slope*xs - ys)**2)

def least_squares_through_point(*, xs: np.ndarray, ys: np.ndarray, required_x: float, required_y: float) -> LinregressResult:
    xs2 = xs - required_x
    ys2 = ys - required_y

    def err(slope: float) -> float:
        return least_squares_cost(xs=xs2, ys=ys2, intercept=0, slope=slope)

    (best_slope,), _ = leastsq(func=err, x0=0.0)
    intercept = required_y - required_x * best_slope
    return LinregressResult(best_slope, intercept, None, None, None, intercept_stderr=False)

def least_squares_with_slope(*, xs: np.ndarray, ys: np.ndarray, required_slope: float) -> LinregressResult:
    def err(intercept: float) -> float:
        return least_squares_cost(xs=xs, ys=ys, intercept=intercept, slope=required_slope)

    (best_intercept,), _ = leastsq(func=err, x0=0.0)
    return LinregressResult(required_slope, best_intercept, None, None, None, intercept_stderr=False)

def least_squares_output_range(*,
                               xs: Sequence[float],
                               ys: Sequence[float],
                               target_x: float,
                               cost_increase: float) -> Tuple[float, float, float]:
    xs = np.array(xs, dtype=np.float64)
    ys = np.array(ys, dtype=np.float64)
    fit = linregress(xs, ys)
    base_cost = least_squares_cost(xs=xs, ys=ys, intercept=fit.intercept, slope=fit.slope)
    base_y = float(fit.intercept + target_x * fit.slope)

    def cost_for_y(y2: float) -> float:
        fit2 = least_squares_through_point(xs=xs, ys=ys, required_x=target_x, required_y=y2)
        return least_squares_cost(xs=xs, ys=ys, intercept=fit2.intercept, slope=fit2.slope)

    low_y = binary_intercept(start_x=base_y, step=-1, target_y=base_cost + cost_increase, func=cost_for_y, atol=1e-5)
    high_y = binary_intercept(start_x=base_y, step=1, target_y=base_cost + cost_increase, func=cost_for_y, atol=1e-5)
    return low_y, base_y, high_y

def least_squares_slope_range(*,
                              xs: Sequence[float],
                              ys: Sequence[float],
                              cost_increase: float) -> Tuple[float, float, float]:
    xs = np.array(xs, dtype=np.float64)
    ys = np.array(ys, dtype=np.float64)
    fit = linregress(xs, ys)
    base_cost = least_squares_cost(xs=xs, ys=ys, intercept=fit.intercept, slope=fit.slope)

    def cost_for_slope(slope: float) -> float:
        fit2 = least_squares_with_slope(xs=xs, ys=ys, required_slope=slope)
        return least_squares_cost(xs=xs, ys=ys, intercept=fit2.intercept, slope=fit2.slope)

    low_slope = binary_intercept(start_x=fit.slope, step=-1, target_y=base_cost + cost_increase, func=cost_for_slope, atol=1e-5)
    high_slope = binary_intercept(start_x=fit.slope, step=1, target_y=base_cost + cost_increase, func=cost_for_slope, atol=1e-5)
    return low_slope, fit.slope, high_slope


from sklearn.ensemble import GradientBoostingClassifier
from commons import variables
from commons import tools
from scipy.stats import mode

def learn(x, y, test_x):
    # set sample weight

    weight_list = []
    for j in range(len(y)):
        if y[j] == "0":
            weight_list.append(variables.weight_0_gdbt_b)
        if y[j] == "1000":
            weight_list.append(variables.weight_1000_gdbt_b)
        if y[j] == "1500":
            weight_list.append(variables.weight_1500_gdbt_b)
        if y[j] == "2000":
            weight_list.append(variables.weight_2000_gdbt_b)

    clf = GradientBoostingClassifier(loss='deviance', n_estimators=variables.n_estimators_gdbt_b,
                                     learning_rate=variables.learning_rate_gdbt_b,
                                     max_depth=variables.max_depth_gdbt_b, random_state=0,
                                     min_samples_split=variables.min_samples_split_gdbt_b,
                                     min_samples_leaf=variables.min_samples_leaf_gdbt_b,
                                     subsample=variables.subsample_gdbt_b,
                                     ).fit(x, y, weight_list)
    prediction_list = clf.predict(test_x)

    return prediction_list


'''original example for checking how far GAM works

Note: uncomment plt.show() to display graphs
'''

example = 2  # 1,2 or 3

import numpy as np
import numpy.random as R
import matplotlib.pyplot as plt

from statsmodels.sandbox.gam import AdditiveModel
from statsmodels.sandbox.gam import Model as GAM #?
from statsmodels.genmod.families import family
from statsmodels.genmod.generalized_linear_model import GLM

standardize = lambda x: (x - x.mean()) / x.std()
demean = lambda x: (x - x.mean())
nobs = 150
x1 = R.standard_normal(nobs)
x1.sort()
x2 = R.standard_normal(nobs)
x2.sort()
y = R.standard_normal((nobs,))

f1 = lambda x1: (x1 + x1**2 - 3 - 1 * x1**3 + 0.1 * np.exp(-x1/4.))
f2 = lambda x2: (x2 + x2**2 - 0.1 * np.exp(x2/4.))
z = standardize(f1(x1)) + standardize(f2(x2))
z = standardize(z) * 2 # 0.1

y += z
d = np.array([x1,x2]).T

if example == 1:
    print "normal"
    m = AdditiveModel(d)
    m.fit(y)
    x = np.linspace(-2,2,50)

    print m

    y_pred = m.results.predict(d)
    plt.figure()
    plt.plot(y, '.')
    plt.plot(z, 'b-', label='true')
    plt.plot(y_pred, 'r-', label='AdditiveModel')
    plt.legend()
    plt.title('gam.AdditiveModel')

import scipy.stats, time

if example == 2:
    print "binomial"
    f = family.Binomial()
    b = np.asarray([scipy.stats.bernoulli.rvs(p) for p in f.link.inverse(y)])
    b.shape = y.shape
    m = GAM(b, d, family=f)
    toc = time.time()
    m.fit(b)
    tic = time.time()
    print tic-toc

if example == 3:
    print "Poisson"
    f = family.Poisson()
    y = y/y.max() * 3
    yp = f.link.inverse(y)
    p = np.asarray([scipy.stats.poisson.rvs(p) for p in f.link.inverse(y)], float)
    p.shape = y.shape
    m = GAM(p, d, family=f)
    toc = time.time()
    m.fit(p)
    tic = time.time()
    print tic-toc

plt.figure()
plt.plot(x1, standardize(m.smoothers[0](x1)), 'r')
plt.plot(x1, standardize(f1(x1)), linewidth=2)
plt.figure()
plt.plot(x2, standardize(m.smoothers[1](x2)), 'r')
plt.plot(x2, standardize(f2(x2)), linewidth=2)

plt.show()

##     pylab.figure(num=1)
##     pylab.plot(x1, standardize(m.smoothers[0](x1)), 'b')
##     pylab.plot(x1, standardize(f1(x1)), linewidth=2)
##     pylab.figure(num=2)
##     pylab.plot(x2, standardize(m.smoothers[1](x2)), 'b')
##     pylab.plot(x2, standardize(f2(x2)), linewidth=2)
##     pylab.show()


<gh_stars>0
import math
import os
import random

import bpy
import scipy.misc

IMAGE_SIZE = 64
DISTANCE = 2.732

def set_camera_location(elevation, azimuth, distance):
    # set location
    x = 1 * math.cos(math.radians(-azimuth)) * math.cos(math.radians(elevation)) * distance
    y = 1 * math.sin(math.radians(-azimuth)) * math.cos(math.radians(elevation)) * distance
    z = 1 * math.sin(math.radians(elevation)) * distance
    camera = bpy.data.objects["Camera"]
    camera.location = x, y, z

    # look at center
    direction = - camera.location
    rot_quat = direction.to_track_quat('-Z', 'Y')
    camera.rotation_euler = rot_quat.to_euler()

def render(directory, elevation=30, distance=DISTANCE):
    for azimuth in range(0, 360, 15):
        filename = os.path.join(directory, 'e%03d_a%03d.png' % (elevation, azimuth))
        set_camera_location(elevation, azimuth, distance)
        bpy.context.scene.render.filepath = filename
        bpy.ops.render.render(write_still=True)

        if False:
            img = scipy.misc.imread(filename)[:, :, :].astype('float32') / 255.
            if False:
                img = (img[::2, ::2] + img[1::2, ::2] + img[::2, 1::2] + img[1::2, 1::2]) / 4.
            else:
                import chainer.functions as cf
                img = img.transpose((2, 0, 1))[None, :, :, :]
                img = cf.resize_images(img, (64, 64))
                img = img[0].data.transpose((1, 2, 0))

            img = (img * 255).clip(0., 255.).astype('uint8')
            scipy.misc.imsave(filename, img)

def setup():
    context = bpy.context
    if False:
        context.scene.render.resolution_x = IMAGE_SIZE * 2
        context.scene.render.resolution_y = IMAGE_SIZE * 2
        context.scene.render.resolution_percentage = 100
        context.scene.render.use_antialiasing = False
    else:
        context.scene.render.resolution_x = IMAGE_SIZE
        context.scene.render.resolution_y = IMAGE_SIZE
        context.scene.render.resolution_percentage = 100
        context.scene.render.use_antialiasing = True
    context.scene.render.use_free_unused_nodes = True
    context.scene.render.use_free_image_textures = True
    context.scene.render.alpha_mode = 'TRANSPARENT'
    bpy.context.scene.render.image_settings.color_mode = 'RGBA'

    # camera
    camera = bpy.data.cameras.values()[0]
    camera.sensor_width = 1
    camera.sensor_height = 1
    camera.lens = 1.8660254037844388

    # lighting
    light = bpy.data.objects['Lamp']
    light.data.energy = 1
    context.scene.world.light_settings.use_environment_light = True
    context.scene.world.light_settings.environment_energy = 0.5
    context.scene.world.light_settings.environment_color = 'PLAIN'

def load_obj(filename):
    # filename = '/home/hkato/temp/obj/model.obj'
    # filename = '/media/disk2/lab/large_data/ShapeNetCore.v1/03001627/1bcec47c5dc259ea95ca4adb70946a21/model.obj'
    bpy.ops.import_scene.obj(filepath=filename, use_smooth_groups=False, use_split_objects=False,
                             use_split_groups=False)
    object_id = len(bpy.data.objects) - 1
    obj = bpy.data.objects[object_id]
    bpy.context.scene.objects.active = obj

    # get max & min of vertices
    inf = 10000
    vertex_max = [-inf, -inf, -inf]
    vertex_min = [inf, inf, inf]
    for j in range(8):
        for i in range(3):
            vertex_max[i] = max(vertex_max[i], obj.bound_box[j][i])
            vertex_min[i] = min(vertex_min[i], obj.bound_box[j][i])
    dimensions = obj.dimensions  # = max - min

    # centering
    for i in range(3):
        obj.location[i] += (vertex_max[i] + vertex_min[i]) / 2

    # scaling
    scale = max(dimensions)
    for i in range(3):
        obj.scale[i] = obj.scale[i] / scale

    # materials
    for m in bpy.data.materials:
        m.ambient = 0.5
        m.use_shadeless = False
        m.use_transparency = False
        m.use_raytrace = False

def clear():
    bpy.ops.wm.open_mainfile(filepath='/home/hkato/temp/untitled.blend')

def run():

    #se crea una lista de id, dentro del dataset para entrenar el modelo con esos objetos
    """class_ids = [
                    '02691156', '02828884', '02933112', '02958343', '03001627', '03211117', '03636649', '03691459', '04090263',
                    '04256520', '04379243', '04401088', '04530566']"""
    class_ids = ['02942699']

    """
    directory_shapenet_id = '../../resource/shapenetcore_ids'
    directory_rendering = '/media/disk2/lab/projection/reconstruction/shapenet_images_%d_%.1f/%s/%s'
    filename_shapenet_obj = '/media/disk2/lab/large_data/ShapeNetCore.v1/%s/%s/model.obj'
    """

    """
    directory_shapenet_id = '/media/gustavo/gusgus/TROOPS/shapenetcore_ids'
    directory_rendering = '/media/gustavo/gusgus/TROOPS/mesh_reconstruction/shapenet_images_%d_%.1f/%s/%s'
    filename_shapenet_obj = '/media/gustavo/gusgus/TROOPS/ShapeNetCore.v2/%s/%s/model.obj'
    """

    directory_shapenet_id = '/content/3D-Sneakers/shapenetcore_ids'
    directory_rendering = '/content/3D-Sneakers/mesh_reconstruction/shapenet_images_%d_%.1f/%s/%s'
    filename_shapenet_obj = '/content/3D-Sneakers/ShapeNetCore.v2/%s/%s/model.obj'

    # ce33bf3ec6438e5bef662d1962a11f02
    for class_id in class_ids:

        ids = open(os.path.join(directory_shapenet_id, '%s_trainids.txt' % class_id), 'w+').readlines()
        ids += open(os.path.join(directory_shapenet_id, '%s_valids.txt' % class_id), 'w+').readlines()
        ids += open(os.path.join(directory_shapenet_id, '%s_testids.txt' % class_id), 'w+').readlines()
        obj_ids = [i.strip().split('/')[-1] for i in ids if len(i.strip()) != 0]

        for i, obj_id in enumerate(obj_ids):
            print('rendering: %s %d / %d' % (class_id, i, len(obj_ids)))

            directory = directory_rendering % (IMAGE_SIZE, DISTANCE, class_id, obj_id)
            directory_tmp = directory + '_'
            if os.path.exists(directory):
                continue
            if os.path.exists(directory_tmp):
                continue
            try:
                os.makedirs(directory_tmp)
            except:
                continue

            clear()
            setup()
            load_obj(filename_shapenet_obj % (class_id, obj_id))
            render(directory_tmp)
            try:
                os.rename(directory_tmp, directory)
            except:
                continue

run()


# -*- coding: utf-8 -*-

"""Main module."""
import numpy as np
import scipy.stats as ss
import matplotlib.pyplot as plt
import argparse

def getinput():
    input = argparse.ArgumentParser()
    input.add_argument('--initial_position', type = float, default = 0, help = 'Initial position of the particle, default = 0' )
    input.add_argument('--initial_velocity', type = float, default = 0, help = 'Initial velocity of the particle, default = 0' )
    input.add_argument('--temperature', type = float, default = 300, help = 'Temperature of the molecule, default = 300' )
    input.add_argument('--damping_coefficient', type = float, default = 0.1, help = 'Damping Coefficient of the molecule, default = 0.1' )
    input.add_argument('--time_step', type = float, default = 0.2, help = 'Time interval of the simulation, default = 0.01' )
    input.add_argument('--wall_size', type = float, default = 5, help = 'Wall size of the simulation, default = 5' )
    input.add_argument('--total_time', type = float, default = 1000, help = 'Total time of the simulation, default = 1000' )
    inp = input.parse_args()
    return inp

def acceleration(gamma=0.1,velocity=0,temperature=300,timestep=0.1,mass=1):
    sigma=np.sqrt(2*temperature*gamma*1*timestep)
    return (-gamma*velocity/mass + np.random.normal(0,sigma))*timestep

def checkwall(position, wallsize):
    if position >= wallsize or position<=0:
        return True
    else:
        return False

def lgmotion(velocity,timestep):
    return velocity*timestep

def integrate(position=0,velocity=0,temperature=300,gamma=0.1,timestep=0.1,wallsize=5,totaltime=1000,mass=1):

    timepass=0
    indexnum=0
    index=[]

    while timepass < totaltime :
        indexnum +=1
        index.append([indexnum,timepass,position,velocity])
        timepass+=timestep
        velocity += acceleration(gamma, velocity, temperature, timestep)
        position += lgmotion(velocity, timestep)
        if checkwall(position,wallsize):
            if position >= wallsize:
                position = wallsize
                index.append([indexnum+1,timepass,position,velocity])
            else:
                position= 0
                index.append([indexnum+1,timepass,position,velocity])
            break

    return timepass,index

def filecreation(index):
    indexf=np.array(index)
    timef=indexf[:,1]
    positionf=indexf[:,2]
    velocityf=indexf[:,3]
    with open('Langevin_Motion.txt','w+') as file:
        file.write('Index Time Position Velocity \n')
        for i in range(len(timef)):
            file.write('{}  {:.3f}  {:.5f}  {:.5f} \n'.format(i,timef[i],positionf[i],velocityf[i]))

def histogram(arr):
    plt.figure(0)
    plt.hist(arr,bins=20)
    plt.title('100 runs of Langevin Motion')
    plt.xlabel('Time passed')
    plt.ylabel('Number of runs')
    plt.savefig('histogram.png')

def trajectory(x,y):
    plt.figure(1)
    plt.plot(x,y)
    plt.title('Position vs Time')
    plt.xlabel('Time passed')
    plt.ylabel('Position')
    plt.savefig('trajectory.png')

def main():
    #get input for simulation
    inp=getinput()

    #run for 100 times, collecting all the relavant data

    t_arr=[] #time
    for i in range(100):
        t,idx=integrate(position=inp.initial_position,velocity=inp.initial_velocity,temperature=inp.temperature,gamma=inp.damping_coefficient,timestep=inp.time_step,wallsize=inp.wall_size,totaltime=inp.total_time,mass=1)
        t_arr.append(t)

    #plot the histogram of 100 runs
    histogram(t_arr)

    #plot the position vs time plot of the last run
    trjdata=np.array(idx)
    xdata=trjdata[:,1]
    ydata=trjdata[:,2]
    trajectory(xdata,ydata)

    #write the index in to a txt file of the first run
    filecreation(idx)

if __name__ == '__main__':
    main()


<gh_stars>0
import numpy as np
from scipy import interpolate
import os
import sys
import subprocess
import abc
from six import with_metaclass

try:
    devnull = open(os.devnull, 'w')
    subprocess.call(['gnuplot', '--version'], stdout=devnull, stderr=devnull)
    import gnuplotpy as gp
    MPL = False
except:
    import matplotlib.pylab as plt
    MPL = True

def use_gnuplot():
    """
    Use gnuplot as the plotting tool for any structure related outputs.
    """
    global gp
    import gnuplotpy as gp
    global MPL
    MPL = False

def use_matplotlib():
    """
    Use matplotlib as the plotting tool for any structure related outputs.
    """
    global plt
    import matplotlib.pylab as plt
    global MPL
    MPL = True

class _AbstractStructure(with_metaclass(abc.ABCMeta)):
    @abc.abstractproperty
    def n(self):
        '''
        np.array: A grid of refractive indices representing
        the refractive index profile of the structure.
        '''
        pass

    @property
    def x_pts(self):
        '''
        int: The number of grid points in x.
        '''
        return int((self.x_max - self.x_min) / self.x_step + 1)

    @property
    def y_pts(self):
        '''
        int: The number of grid points in y.
        '''
        return int((self.y_max - self.y_min) / self.y_step)

    @property
    def x_ctr(self):
        '''
        float: The centre distance in x.
        '''
        return 0.5*(self.x_max + self.x_min)

    @property
    def y_ctr(self):
        '''
        float: The centre distance in y
        '''
        return 0.5*(self.y_max + self.y_min)

    @property
    def xc(self):
        '''
        np.array: The centre points of the x points.
        '''
        return 0.5*(self.x[1:] + self.x[:-1])

    @property
    def yc(self):
        '''
        np.array: The centre points of the y points.
        '''
        return 0.5*(self.y[1:] + self.y[:-1])

    @property
    def xc_pts(self):
        '''
        int: The number of points in `xc`.
        '''
        return self.x_pts - 1

    @property
    def yc_pts(self):
        '''
        int: The number of points in `yc`.
        '''
        return self.y_pts - 1

    @property
    def xc_min(self):
        '''
        float: The minimum value of `xc`.
        '''
        return self.xc[0]

    @property
    def xc_max(self):
        '''
        float: The maximum value of `xc`.
        '''
        return self.xc[-1]

    @property
    def yc_min(self):
        '''
        float: The minimum value of `yc`.
        '''
        return self.yc[0]

    @property
    def yc_max(self):
        '''
        float: The maximum value of `yc`.
        '''
        return self.yc[-1]

    @property
    def x(self):
        '''
        np.array: The grid points in x.
        '''
        if None not in (self.x_min, self.x_max, self.x_step) and \
                self.x_min != self.x_max:
            x = np.arange(self.x_min, self.x_max+self.x_step-self.y_step*0.1, self.x_step)
        else:
            x = np.array([])
        return x

    @property
    def y(self):
        '''
        np.array: The grid points in y.
        '''
        if None not in (self.y_min, self.y_max, self.y_step) and \
                self.y_min != self.y_max:
            y = np.arange(self.y_min, self.y_max-self.y_step*0.1, self.y_step)
        else:
            y = np.array([])
        return y

    @property
    def eps(self):
        '''
        np.array: A grid of permittivies representing
        the permittivity profile of the structure.
        '''
        return self.n**2

    @property
    def eps_func(self):
        '''
        function: a function that when passed a `x` and `y` values,
            returns the permittivity profile of the structure,
            interpolating if necessary.
        '''
        interp_real = interpolate.interp2d(self.x, self.y, self.eps.real)
        interp_imag = interpolate.interp2d(self.x, self.y, self.eps.imag)
        interp = lambda x, y: interp_real(x, y) + 1.j*interp_imag(x, y)
        return interp

    @property
    def n_func(self):
        '''
        function: a function that when passed a `x` and `y` values,
            returns the refractive index profile of the structure,
            interpolating if necessary.
        '''
        return interpolate.interp2d(self.x, self.y, self.n)

    def _add_triangular_sides(self, xy_mask, angle, y_top_right, y_bot_left,
                              x_top_right, x_bot_left, n_material):
        angle = np.radians(angle)
        trap_len = (y_top_right - y_bot_left) / np.tan(angle)
        num_x_iterations = trap_len / self.x_step
        y_per_iteration =  num_x_iterations / self.y_pts

        lhs_x_start_index = int(x_bot_left/ self.x_step + 0.5)
        rhs_x_stop_index = int(x_top_right/ self.x_step + 1 + 0.5)

        running_removal_float = y_per_iteration
        for i, _ in enumerate(xy_mask):
            if running_removal_float >= 1:
                removal_int = int(round(running_removal_float))
                lhs_x_start_index -= removal_int
                rhs_x_stop_index += removal_int
                running_removal_float -= removal_int
            running_removal_float += y_per_iteration

            xy_mask[i][:lhs_x_start_index] = False
            xy_mask[i][lhs_x_start_index:rhs_x_stop_index] = True

        self.n[xy_mask] = n_material
        return self.n

    def _add_material(self, x_bot_left, y_bot_left, x_top_right, y_top_right,
                     n_material, angle=0):
        '''
        A low-level function that allows writing a rectangle refractive
        index profile to a `Structure`.

        Args:
            x_bot_left (float): The bottom-left x-coordinate of the
                rectangle.
            y_bot_left (float): The bottom-left y-coordinate of the
                rectangle.
            x_top_right (float): The top-right x-coordinate of the
                rectangle.
            y_top_right (float): The top-right y-coordinate of the
                rectangle.
            n_material (float): The refractive index of the points
                encompassed by the defined rectangle.
            angle (float): The angle in degrees of the sidewalls
                of the defined rectangle.  Default is 0.  This
                is useful for creating a ridge with angled
                sidewalls.
        '''
        x_mask = np.logical_and(x_bot_left<=self.x, self.x<=x_top_right)
        y_mask = np.logical_and(y_bot_left<=self.y, self.y<=y_top_right)

        xy_mask = np.kron(y_mask, x_mask).reshape((y_mask.size, x_mask.size))
        self.n[xy_mask] = n_material

        if angle:
            self._add_triangular_sides(xy_mask, angle, y_top_right, y_bot_left,
                                       x_top_right, x_bot_left, n_material)

        return self.n

    def write_to_file(self, filename='material_index.dat', plot=True):
        '''
        Write the refractive index profile to file.

        Args:
            filename (str): The nominal filename the refractive
                index data should be saved to.
            plot (bool): `True` if plots should be generates,
                otherwise `False`.  Default is `True`.
        '''
        path = os.path.dirname(sys.modules[__name__].__file__) + '/'

        with open(filename, 'w') as fs:
            for n_row in np.abs(self.n[::-1]):
                n_str = ','.join([str(v) for v in n_row])
                fs.write(n_str+'\n')

        if plot:
            filename_image_prefix, _ = os.path.splitext(filename)
            filename_image = filename_image_prefix + '.png'
            args = {
                'title': 'Refractive Index Profile',
                'x_pts': self.x_pts,
                'y_pts': self.y_pts,
                'x_min': self.x_min,
                'x_max': self.x_max,
                'y_min': self.y_min,
                'y_max': self.y_max,
                'filename_data': filename,
                'filename_image': filename_image
            }

            if MPL:
                heatmap = np.loadtxt(args['filename_data'], delimiter=',')
                structure_plot = plt.figure()
                plt.clf()
                plt.title(args['title'])
                plt.xlabel('$x$')
                plt.ylabel('$y$')
                plt.imshow(np.flipud(heatmap),
                           extent=(args['x_min'], args['x_max'], args['y_min'], args['y_max']),
                           aspect="auto")
                plt.colorbar()
                plt.savefig(filename_image)
                #plt.show()
                print("Returning structure plot")
                return structure_plot
            else:
                print("using gnu plot")
                gp.gnuplot(path+'structure.gpi', args)

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

class Structure(_AbstractStructure):
    def __init__(self, x_step, y_step, x_max, y_max, x_min=0., y_min=0.,
                 n_background=1.):
        self.x_min = x_min
        self.x_max = x_max
        self.y_min = y_min
        self.y_max = y_max
        self.x_step = x_step
        self.y_step = y_step
        self.n_background = n_background
        self._n = np.ones((self.y.size,self.x.size), 'complex_') * n_background

    @property
    def n(self):
        return self._n

class Slabs(_AbstractStructure):
    '''
    Class to implement device refractive index
    profile cross-section designs.

    :class:`Slabs` is a collection of :class:`Slab` objects.  Each
    slab has a fixed height (usually less than the
    maximum height of the desired simulation window),
    and is as wide as the simulation window.

    :class:`Slabs` objects can be index using `[name]` to return
    the various :class:`Slab` objects.  The bottom slab is
    returned first and so on up to the top slab.

    .. image:: ../images/slabs.svg
        :width: 200%

    Args:
        wavelength (float): The wavelength the structure
            operates at.
        y_step (float): The step in y.
        x_step (float): The step in x.
        x_max (float): The maximum x-value.
        x_min (float): The minimum x-value. Default is 0.

    Attributes:
        slabs (dict): The key is the name of the slab,
            and the value is the :class:`Slab` object.
        slab_count (int): The number of :class:`Slab` objects
            added so far.
    '''
    def __init__(self, wavelength, y_step, x_step, x_max, x_min=0.):
        _AbstractStructure.__init__(self)

        self._wl = wavelength
        self.x_min = x_min
        self.x_max = x_max
        self.x_step = x_step
        self.y_step = y_step
        self.y_min = 0

        self.slabs = {}
        self.slab_count = 0
        self._next_start = 0.

    def add_slab(self, height, n_background=1., position='top'):
        '''
        Creates and adds a :class:`Slab` object.

        Args:
            height (float): Height of the slab.
            n_background (float): The nominal refractive
                index of the slab.  Default is 1 (air).

        Returns:
            str: The name of the slab.
        '''
        assert position in ('top', 'bottom')

        name = str(self.slab_count)

        if not callable(n_background):
            n_back = lambda wl: n_background
        else:
            n_back = n_background

        height_discretised = self.y_step*((height // self.y_step) + 1)

        y_min = self._next_start
        y_max = y_min + height_discretised
        self.slabs[name] = Slab(name, self.x_step, self.y_step, self.x_max,
                                 y_max, self.x_min, y_min, n_back, self._wl)

        self.y_max = y_max
        self._next_start = y_min + height_discretised
        self.slab_count += 1

        if position == 'bottom':
            slabs = {}
            for k in self.slabs.keys():
                slabs[str(int(k)+1)] = self.slabs[k]
            slabs['0'] = slabs.pop(str(self.slab_count))
            self.slabs = slabs

        return name

    def change_wavelength(self, wavelength):
        '''
        Changes the wavelength of the structure.

        This will affect the mode solver and potentially
        the refractive indices used (provided functions
        were provided as refractive indices).

        Args:
            wavelength (float): The new wavelength.
        '''
        for name, slab in self.slabs.items():
            const_args = slab._const_args
            mat_args = slab._mat_params

            const_args[8] = wavelength

            s = Slab(*const_args)
            for mat_arg in mat_args:
                s.add_material(*mat_arg)

            self.slabs[name] = s

        self._wl = wavelength

    @property
    def n(self):
        '''
        np.array: The refractive index profile matrix
        of the current slab.
        '''
        try:
            n_mat = self.slabs['0'].n
            for s in range(1, self.slab_count):
                n_mat = np.vstack((self.slabs[str(s)].n, n_mat))
        except KeyError:
            n_mat = None
        return n_mat

    def __getitem__(self, slab_name):
        return self.slabs[str(slab_name)]

class Slab(Structure):
    '''
    A :class:`Slab` represents a horizontal slice of
    the refractive index profile.

    A :class:`Slabs` object composes many :class:`Slab` objects.
    The more :class:`Slab` are added, the more horizontal
    slices are added.  A :class:`Slab` has a chosen fixed
    height, and a background (nominal) refractive
    index.  A slab can then be customised to include
    a desired design.

    Args:
        name (str): The name of the slab.
        x_step (float): The step in x.
        y_step (float): The step in y.
        x_max (float): The maximum x-value.
        y_max (float): The maximum y-value.
        x_min (float): The minimum x-value.
        y_min (float): The minimum x-value.
        n_background (float): The nominal refractive
            index.
        wavelength (float): The wavelength the structure
            operates at.

    Attributes:
        name (str): The name of the :class:`Slab` object.
        position (int): A unique identifier for the
        :class:`Slab` object.
    '''
    position = 0

    def __init__(self, name, x_step, y_step, x_max, y_max, x_min, y_min,
                 n_background, wavelength):
        self._wl = wavelength
        self.name = name
        self.position = Slab.position
        Slab.position += 1

        Structure.__init__(self, x_step, y_step, x_max, y_max, x_min, y_min,
                           n_background(self._wl))

        self._const_args = [name, x_step, y_step, x_max, y_max, x_min, y_min, n_background, wavelength]
        self._mat_params = []

    def add_material(self, x_min, x_max, n, angle=0):
        '''
        Add a refractive index between two x-points.

        Args:
            x_min (float): The start x-point.
            x_max (float): The stop x-point.
            n (float, function):  Refractive index between
                 `x_min` and `x_max`.  Either a constant (`float`), or
                 a function that accepts one parameters, the
                 wavelength, and returns a float of the refractive
                 index.  This is useful when doing wavelength
                 sweeps and solving for the group velocity.  The
                 function provided could be a Sellmeier equation.
            angle (float): Angle in degrees of the slope of the
                sidewalls at `x_min` and `x_max`.  This is useful
                for defining a ridge with angled sidewalls.
        '''
        self._mat_params.append([x_min, x_max, n, angle])

        if not callable(n):
            n_mat = lambda wl: n
        else:
            n_mat = n

        Structure._add_material(self, x_min, self.y_min, x_max, self.y_max, n_mat(self._wl), angle)
        return self.n

class StructureAni():
    r"""
    Anisottropic structure object.

    This is used with the fully-vectorial simulation when
    an anisotropic material is being used.

    The form of the refractive index is

    .. math::

        n = \begin{bmatrix}
                n_{xx} & n_{xy} & 0 \\
                n_{yx} & n_{yy} & 0 \\
                0      & 0      & n_{zz}
            \end{bmatrix}.

    Args:
        structure_xx (Structure): The structure with refractive
            index, :math:`n_{xx}`.
        structure_yy (Structure): The structure with refractive
            index, :math:`n_{yy}`.  Presumably the same structure
            as `structure_xx`, but with different refractive index
            parameters.
        structure_zz (Structure): The structure with refractive
            index, :math:`n_{zz}`.  Presumably the same structure
            as `structure_xx`, but with different refractive index
            parameters.
        structure_xy (None, Structure): The structure with refractive
            index, :math:`n_{yx}`.  Presumably the same structure
            as `structure_xx`, but with different refractive index
            parameters.  Default is `None`.
        structure_yx (None, Structure): The structure with refractive
            index, :math:`n_{yx}`.  Presumably the same structure
            as `structure_xx`, but with different refractive index
            parameters.  Default is `None`.
    """
    def __init__(self, structure_xx, structure_yy, structure_zz,
                 structure_xy=None, structure_yx=None):
        self.xx = structure_xx
        self.yy = structure_yy
        self.zz = structure_zz

        if not structure_xy or not structure_yx:
            struct_dummy = Structure(self.xx.x_step, self.xx.y_step,
                                     self.xx.x_max, self.xx.y_max,
                                     self.xx.x_min, self.xx.y_min,
                                     n_background=0.)
            struct_dummy._wl = self.xx._wl

        if structure_xy:
            self.xy = structure_xy
        else:
            self.xy = struct_dummy

        if structure_yx:
            self.yx = structure_yx
        else:
            self.yx = struct_dummy

        assert self.xx._wl == self.xy._wl == self.yx._wl == \
               self.yy._wl == self.zz._wl

        self._wl = structure_xx._wl

        self.axes = (self.xx, self.xy, self.yx, self.yy, self.zz)
        self.axes_str = ('xx', 'xy', 'yx', 'yy', 'zz')

    @property
    def n(self):
        return [a.n for a in self.axes]

    @property
    def x_step(self):
        return self.xx.x_step

    @property
    def y_step(self):
        return self.xx.y_step

    @property
    def x_pts(self):
        return int((self.xx.x_max - self.xx.x_min) / self.xx.x_step + 1)

    @property
    def y_pts(self):
        return int((self.xx.y_max - self.xx.y_min) / self.xx.y_step)

    @property
    def x_ctr(self):
        return 0.5*(self.xx.x_max + self.xx.x_min)

    @property
    def y_ctr(self):
        return 0.5*(self.xx.y_max + self.xx.y_min)

    @property
    def xc(self):
        return 0.5*(self.xx.x[1:] + self.xx.x[:-1])

    @property
    def yc(self):
        return 0.5*(self.xx.y[1:] + self.xx.y[:-1])

    @property
    def xc_pts(self):
        return self.xx.x_pts - 1

    @property
    def yc_pts(self):
        return self.xx.y_pts - 1

    @property
    def xc_min(self):
        return self.xx.xc[0]

    @property
    def xc_max(self):
        return self.xx.xc[-1]

    @property
    def yc_min(self):
        return self.xx.yc[0]

    @property
    def yc_max(self):
        return self.xx.yc[-1]

    @property
    def x(self):
        if None not in (self.xx.x_min, self.xx.x_max, self.xx.x_step) and \
                self.xx.x_min != self.xx.x_max:
            x = np.arange(self.xx.x_min, self.xx.x_max+self.xx.x_step-self.xx.y_step*0.1, self.xx.x_step)
        else:
            x = np.array([])
        return x

    @property
    def y(self):
        if None not in (self.xx.y_min, self.xx.y_max, self.xx.y_step) and \
                self.xx.y_min != self.xx.y_max:
            y = np.arange(self.xx.y_min, self.xx.y_max-self.xx.y_step*0.1, self.xx.y_step)
        else:
            y = np.array([])
        return y

    @property
    def eps(self):
        eps_ani = [a.n**2 for a in self.axes]
        return eps_ani

    @property
    def eps_func(self):
        return lambda x,y: tuple(axis.eps_func(x,y) for axis in self.axes)

    @property
    def n_func(self):
        return lambda x,y: tuple(axis.n_func(x,y) for axis in self.axes)

    def write_to_file(self, filename='material_index.dat', plot=True):
        '''
        Write the refractive index profile to file.

        Args:
            filename (str): The nominal filename the refractive
                index data should be saved to.
            plot (bool): `True` if plots should be generates,
                otherwise `False`.  Default is `True`.
        '''
        path = os.path.dirname(sys.modules[__name__].__file__) + '/'

        dir_plot = 'material_index/'
        if not os.path.exists(dir_plot):
            os.makedirs(dir_plot)

        for axis, name in zip(self.axes, self.axes_str):
            root, ext = os.path.splitext(filename)
            fn = dir_plot + root + '_'+ name + ext
            with open(fn, 'w') as fs:
                for n_row in np.abs(axis.n[::-1]):
                    n_str = ','.join([str(v) for v in n_row])
                    fs.write(n_str+'\n')

            if plot:
                filename_image_prefix, _ = os.path.splitext(fn)
                filename_image = filename_image_prefix + '.png'
                args = {
                    'title': 'Refractive Index Profile: %s' % name,
                    'x_pts': self.xx.x_pts,
                    'y_pts': self.xx.y_pts,
                    'x_min': self.xx.x_min,
                    'x_max': self.xx.x_max,
                    'y_min': self.xx.y_min,
                    'y_max': self.xx.y_max,
                    'filename_data': fn,
                    'filename_image': filename_image
                }
                if MPL:
                    heatmap = np.loadtxt(args['filename_data'], delimiter=',')
                    plt.clf()
                    plt.title(args['title'])
                    plt.xlabel('$x$')
                    plt.ylabel('$y$')
                    plt.imshow(np.flipud(heatmap),
                               extent=(args['x_min'], args['x_max'], args['y_min'], args['y_max']),
                               aspect="auto")
                    plt.colorbar()
                    plt.savefig(filename_image)
                else:
                    gp.gnuplot(path+'structure.gpi', args, silent=False)

    def change_wavelength(self, wavelength):
        '''
        Changes the wavelength of the structure.

        This will affect the mode solver and potentially
        the refractive indices used (provided functions
        were provided as refractive indices).

        Args:
            wavelength (float): The new wavelength.
        '''
        for axis in self.axes:
            if issubclass(type(axis), Slabs):
                axis.change_wavelength(wavelength)
        self.xx, self.xy, self.yx, self.yy, self.zz = self.axes
        self._wl = wavelength


# -*- coding: utf-8 -*-
import numpy as np
import cvxpy as cp
import random
import time
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KernelDensity
from sklearn.mixture import GaussianMixture
from sklearn.model_selection import GridSearchCV, KFold
from sklearn.decomposition import PCA
from sklearn.datasets import load_iris, load_digits, load_breast_cancer, load_boston, load_wine
from sklearn.metrics import accuracy_score
from sklearn.utils import shuffle

from plausible_counterfactuals import HighDensityEllipsoids, PlausibleCounterfactualOfHyperplaneClassifier, PlausibleCounterfactualOfDecisionTree

def load_house_prices(file_path="housepricesdataset.npz"):
    X, y = load_boston(return_X_y=True)
    y = y >= 20
    y = y.astype(np.int).flatten()

    return X, y

if __name__ == "__main__":
    use_decision_tree = False   # If False, softmax regression is used!

    # Load data set
    X, y = load_iris(return_X_y=True);pca_dim=None
    #X, y = load_breast_cancer(return_X_y=True);pca_dim=5
    #X, y = load_house_prices();pca_dim=10
    #X, y = load_wine(return_X_y=True);pca_dim=8
    X, y = load_digits(return_X_y=True);pca_dim=40

    X, y = shuffle(X, y, random_state=42)

    # k-fold cross validation
    scores_with_density_constraint = []
    scores_without_density_constraint = []

    original_data = []
    original_data_labels = []
    cfs_with_density_constraint = []
    cfs_without_density_constraint = []
    cfs_target_label = []
    computation_time_without_density_constraint = []
    computation_time_with_density_constraint = []
    distances_with_density_constraint = []
    distances_without_density_constraint = []

    kf = KFold(n_splits=5, random_state=42)
    for train_index, test_index in kf.split(X):
        # Split data into training and test set
        X_train, X_test = X[train_index], X[test_index]
        y_train, y_test = y[train_index], y[test_index]

        # Choose target labels
        y_test_target = []
        labels = np.unique(y)
        for i in range(X_test.shape[0]):
            y_test_target.append(random.choice(list(filter(lambda l: l != y_test[i], labels))))
        y_test_target = np.array(y_test_target)

        # If requested: Reduce dimensionality
        X_train_orig = np.copy(X_train)
        X_test_orig = np.copy(X_test)
        projection_matrix = None
        projection_mean_sub = None
        pca = None
        if pca_dim is not None:
            pca = PCA(n_components=pca_dim)
            pca.fit(X_train)

            projection_matrix = pca.components_ # Projection matrix
            projection_mean_sub = pca.mean_

            X_train = np.dot(X_train - projection_mean_sub, projection_matrix.T)
            X_test = np.dot(X_test - projection_mean_sub, projection_matrix.T)

        # Fit classifier
        model = LogisticRegression(multi_class="multinomial", solver="lbfgs", random_state=42)
        if use_decision_tree is True:
            model = DecisionTreeClassifier(max_depth=7, random_state=42)
        model.fit(X_train, y_train)

        # Compute accuracy on test set
        print("Accuracy: {0}".format(accuracy_score(y_test, model.predict(X_test))))

        # For each class, fit density estimators
        density_estimators = {}
        kernel_density_estimators = {}
        labels = np.unique(y)
        for label in labels:
            # Get all samples with the 'correct' label
            idx = y_train == label
            X_ = X_train[idx, :]

            # Optimize hyperparameters
            cv = GridSearchCV(estimator=KernelDensity(), iid=False, param_grid={'bandwidth': np.arange(0.1, 10.0, 0.05)}, n_jobs=-1, cv=5)
            cv.fit(X_)
            bandwidth = cv.best_params_["bandwidth"]
            print("bandwidth: {0}".format(bandwidth))

            cv = GridSearchCV(estimator=GaussianMixture(covariance_type='full'), iid=False, param_grid={'n_components': range(2, 10)}, n_jobs=-1, cv=5)
            cv.fit(X_)
            n_components = cv.best_params_["n_components"]
            print("n_components: {0}".format(n_components))

            # Build density estimators
            kde = KernelDensity(bandwidth=bandwidth)
            kde.fit(X_)

            de = GaussianMixture(n_components=n_components, covariance_type='full', random_state=42)
            de.fit(X_)

            density_estimators[label] = de
            kernel_density_estimators[label] = kde

        # For each point in the test set
        # Compute and plot counterfactual without density constraints
        print("n_test_samples: {0}".format(X_test.shape[0]))
        for i in range(X_test.shape[0]):
            x_orig = X_test[i,:]
            x_orig_orig = X_test_orig[i,:]
            y_orig = y_test[i]
            y_target = y_test_target[i]

            if(model.predict([x_orig]) == y_target):  # Model already predicts target label!
                print("Requested prediction already satisfied")
                continue

            # Compute and plot counterfactual WITH kernel density constraints
            idx = y_train == y_target
            X_ = X_train[idx, :]

            # Build density estimator
            de = density_estimators[y_target]
            kde = kernel_density_estimators[y_target]

            # Compute media NLL of training samples
            # TODO: Move this to the outer loop
            from scipy.stats import multivariate_normal
            densities_training_samples = []
            densities_training_samples_ex = []
            for j in range(X_.shape[0]):
                x = X_[j,:]
                z = []
                dim = x.shape[0]
                for i in range(de.weights_.shape[0]):
                    x_i = de.means_[i]
                    w_i = de.weights_[i]
                    cov = de.covariances_[i]
                    cov = np.linalg.inv(cov)

                    b = -2.*np.log(w_i) + dim*np.log(2.*np.pi) - np.log(np.linalg.det(cov))
                    z.append(np.dot(x - x_i, np.dot(cov, x - x_i)) + b) # NLL

                densities_training_samples.append(np.min(z))
                densities_training_samples_ex.append(z)

            densities_training_samples = np.array(densities_training_samples)
            densities_training_samples_ex = np.array(densities_training_samples_ex)

            # Compute soft cluster assignments
            cluster_prob_ = de.predict_proba(X_)
            density_threshold = np.median(densities_training_samples)
            # Compute high density ellipsoids - constraint: test if sample is included in ellipsoid -> this is the same as the proposed constraint but nummerically much more stable, in particular when we add a dimensionality reduction from a high dimensional space to a low dimensional space
            r = HighDensityEllipsoids(X_, densities_training_samples_ex, cluster_prob_, de.means_, de.covariances_, density_threshold).compute_ellipsoids()

            # Compute counterfactual without any density/plausibility/feasibility constraints
            xcf_t1 = time.time()
            cf = None
            if use_decision_tree is False:
                cf = PlausibleCounterfactualOfHyperplaneClassifier(model.coef_, model.intercept_, n_dims=X_train.shape[1], density_constraint=False, ellipsoids_r=r, gmm_weights=de.weights_, gmm_means=de.means_, gmm_covariances=de.covariances_, projection_matrix=projection_matrix, projection_mean_sub=projection_mean_sub)
            else:
                cf = PlausibleCounterfactualOfDecisionTree(model, n_dims=X_train.shape[1], density_constraint=False, ellipsoids_r=r, gmm_weights=de.weights_, gmm_means=de.means_, gmm_covariances=de.covariances_, projection_matrix=projection_matrix, projection_mean_sub=projection_mean_sub)
            xcf = cf.compute_counterfactual(x_orig_orig, y=y_target)
            xcf_t1 = time.time() - xcf_t1
            if xcf is None:
                print("No counterfactual found!")
                continue

            # Compute counterfactul with proposed density constraint
            xcf_t2 = time.time()
            cf2 = None
            if use_decision_tree is False:
                cf2 = PlausibleCounterfactualOfHyperplaneClassifier(model.coef_, model.intercept_, n_dims=X_train.shape[1], ellipsoids_r=r, gmm_weights=de.weights_, gmm_means=de.means_, gmm_covariances=de.covariances_, projection_matrix=projection_matrix, projection_mean_sub=projection_mean_sub, density_threshold=density_threshold)
            else:
                cf2 = PlausibleCounterfactualOfDecisionTree(model, n_dims=X_train.shape[1], ellipsoids_r=r, gmm_weights=de.weights_, gmm_means=de.means_, gmm_covariances=de.covariances_, projection_matrix=projection_matrix, projection_mean_sub=projection_mean_sub, density_threshold=density_threshold)
            xcf2 = cf2.compute_counterfactual(x_orig_orig, y=y_target)
            xcf_t2 = time.time() - xcf_t2
            if xcf2 is None:
                print("No counterfactual found!")
                continue

            original_data.append(x_orig_orig)
            original_data_labels.append(y_orig)
            cfs_with_density_constraint.append(xcf2)
            cfs_without_density_constraint.append(xcf)
            cfs_target_label.append(y_target)
            computation_time_without_density_constraint.append(xcf_t1)
            computation_time_with_density_constraint.append(xcf_t2)
            distances_with_density_constraint.append(np.sum(np.abs(x_orig_orig - xcf2)))
            distances_without_density_constraint.append(np.sum(np.abs(x_orig_orig - xcf)))

            if pca is not None: # If necessary: Project the counterfactuals to the lower dimensional space where we did the density estimation
                xcf = pca.transform([xcf])
                xcf2 = pca.transform([xcf2])

            # Evaluate
            scores_without_density_constraint.append(kde.score_samples(xcf.reshape(1, -1)))
            scores_with_density_constraint.append(kde.score_samples(xcf2.reshape(1, -1)))

    # Final evaluation
    print("Without density constrain: Median: {0} Mean: {1} Var: {2}".format(np.median(scores_without_density_constraint), np.mean(scores_without_density_constraint), np.var(scores_without_density_constraint)))
    print("With density constrain: Median: {0} Mean: {1} Var: {2}".format(np.median(scores_with_density_constraint), np.mean(scores_with_density_constraint), np.var(scores_with_density_constraint)))

    print("Computation time: With density constraint: {0} Without density constraint: {1}".format(np.median(computation_time_with_density_constraint), np.median(computation_time_without_density_constraint)))
    print("Distances: With density constraint: {0} {1} Without density constraint: {2} {3}".format(np.median(distances_with_density_constraint), np.mean(distances_with_density_constraint), np.median(distances_without_density_constraint), np.mean(distances_without_density_constraint)))

    #"""
    # Plot some samples: Counterfactual generated with vs. without density constraint
    original_data = np.array(original_data)
    original_data_labels = np.array(original_data_labels)
    cfs_with_density_constraint = np.array(cfs_with_density_constraint)
    cfs_without_density_constraint = np.array(cfs_without_density_constraint)
    cfs_target_label = np.array(cfs_target_label)
    np.savez("cfs_comparision_data_softmax_regression", X_original=original_data, y_original=original_data_labels, y_target=cfs_target_label, X_with_density_constraint=cfs_with_density_constraint, X_without_density_constraint=cfs_without_density_constraint)
    #"""


<gh_stars>1-10
import os
import scipy.interpolate as spi

dataLabel = ["one", "two", "three", "four"]

dataRoot = "../toneclassifier/train"

normalLen = 200

for label in dataLabel:
    subsetPath = dataRoot + "/" + label
    dataset = set()
    for filename in os.listdir(subsetPath):
        if filename[0] == ".":
            continue
        if ".engy" in filename:
            dataset.add(filename[0:-5])
        elif ".f0" in filename:
            dataset.add(filename[0:-3])
    try:
        os.makedirs("../data-process-output/trim-interpolation/train/" + label)
    except OSError as err:
        pass

    for dataname in dataset:
        engyfile = open(subsetPath + "/" + dataname + ".engy", "r")
        f0file = open(subsetPath + "/" + dataname + ".f0", "r")
        engy = map(float, engyfile.readlines())
        f0 = map(float, f0file.readlines())
        engyfile.close()
        f0file.close()

        start = None
        end = None

        for i in xrange(len(f0)):
            if (f0[i] > 1e-5):
                start = i
                break
        for i in xrange(len(f0) - 1, -1, -1):
            if (f0[i] > 1e-5):
                end = i + 1
                break
        engy = engy[start:end]
        f0 = f0[start:end]

        dataLen = len(engy)
        k = float(normalLen - 1) / float(dataLen - 1)

        x = [i * k for i in xrange(dataLen)]
        newX = [i * 1.0 for i in xrange(normalLen)]
        newX[-1] = x[-1]
        # tck = spi.splrep(x, engy)
        # newEngy = spi.splev(newX, tck)
        # tck = spi.splrep(x, f0)
        # newF0 = spi.splev(newX, tck)
        func = spi.interp1d(x, engy, kind='cubic')
        newEngy = func(newX)
        func = spi.interp1d(x, f0, kind='cubic')
        newF0 = func(newX)

        engyfile = open("../data-process-output/trim-interpolation/train/" + label + "/" + dataname + ".engy", "w")
        f0file = open("../data-process-output/trim-interpolation/train/" + label + "/" + dataname + ".f0", "w")
        for i in xrange(normalLen):
            engyfile.write("%.5f\n" % newEngy[i])
            f0file.write("%.5f\n" % newF0[i])
        engyfile.close()
        f0file.close()