<jupyter_start><jupyter_text>| Likelihood       ||------------------|| env1, continuous || env2, continuous ||                       | env1, continuous  | env2, continuous  ||-----------------------|-------------------|-------------------|| **env1, continuous**  | random walk, env1 | uniform           || **env2, continuous**  | uniform           | random walk, env2 |+ Change text inputs for continuous_transition_types to classes<jupyter_code>from replay_trajectory_classification.bins import (atleast_2d, get_centers,
                                                   get_grid, get_track_grid,
                                                   get_track_interior)
from dataclasses import dataclass
import networkx as nx
import matplotlib.pyplot as plt

@dataclass
class Environment:
    environment_name: str = None
    place_bin_size: float = 2.0
    track_graph: nx.Graph = None
    edge_order: tuple = None
    edge_spacing: tuple = None
    is_track_interior: np.ndarray = None
    position_range: np.ndarray = None
    infer_track_interior: bool = True
    
    def fit_place_grid(self, position=None, infer_track_interior=True):
        if self.track_graph is None:
            (self.edges_,
             self.place_bin_edges_,
             self.place_bin_centers_,
             self.centers_shape_
            ) = get_grid(position, self.place_bin_size, self.position_range,
                         self.infer_track_interior)

            self.infer_track_interior = infer_track_interior

            if self.is_track_interior is None and self.infer_track_interior:
                self.is_track_interior_ = get_track_interior(
                    position, self.edges_)
            elif is_track_interior is None and not self.infer_track_interior:
                self.is_track_interior_ = np.ones(
                    self.centers_shape_, dtype=np.bool)
        else:
            (
                self.place_bin_centers_,
                self.place_bin_edges_,
                self.is_track_interior_,
                self.distance_between_nodes_,
                self.centers_shape_,
                self.edges_,
                self.track_graph_with_bin_centers_edges_,
                self.original_nodes_df_,
                self.place_bin_edges_nodes_df_,
                self.place_bin_centers_nodes_df_,
                self.nodes_df_
            ) = get_track_grid(self.track_graph, self.edge_order,
                               self.edge_spacing, self.place_bin_size)

        return self
    
    def plot_grid(self, ax=None):
        if self.track_graph is not None:
            if ax is None:
                fig, ax = plt.subplots(figsize=(15, 2))

            plot_graph_as_1D(self.track_graph, self.edge_order,
                             self.edge_spacing, ax=ax)
            for edge in self.edges_[0]:
                ax.axvline(edge.squeeze(), linewidth=0.5, color='black')
            ax.set_ylim((0, 0.1))
        else:
            if ax is None:
                fig, ax = plt.subplots(figsize=(6, 7))
            ax.pcolormesh(self.edges_[0], self.edges_[1], self.is_track_interior_.T,
                          cmap='bone_r')
            ax.set_xticks(self.edges_[0], minor=True)
            ax.set_yticks(self.edges_[1], minor=True)
            ax.grid(True, which='both')


env1 = Environment(environment_name='w_track1',
                  place_bin_size=2.0,
                  track_graph=track_graph1,
                  edge_spacing=edge_spacing,
                  edge_order=edge_order)
env1.fit_place_grid(position_df1.linear_position)

env2 = Environment(environment_name='w_track2',
                   place_bin_size=2.0,
                   track_graph=track_graph2,
                   edge_spacing=edge_spacing,
                   edge_order=edge_order)
env2.fit_place_grid(position_df2.linear_position)

env1, env2
from replay_trajectory_classification.state_transition import _normalize_row_probability, atleast_2d
from scipy.stats import multivariate_normal
import numpy as np

def _random_walk_on_track_graph(
    place_bin_centers, movement_var,
    place_bin_center_ind_to_node,
    distance_between_nodes
):
    state_transition = np.zeros(
        (place_bin_centers.size, place_bin_centers.size))
    gaussian = multivariate_normal(mean=0, cov=movement_var)

    for bin_ind1, node1 in enumerate(place_bin_center_ind_to_node):
        for bin_ind2, node2 in enumerate(place_bin_center_ind_to_node):
            try:
                state_transition[bin_ind1, bin_ind2] = gaussian.pdf(
                    distance_between_nodes[node1][node2]
                )
            except KeyError:
                # bins not on track interior will be -1 and not in distance
                # between nodes
                continue

    return state_transition

@dataclass
class RandomWalk:
    environment_name: str = None
    movement_var: float = 6.0
        
    def make_state_transition(self, environments: dict):
        self.environment = environments[self.environment_name]

        if self.environment.track_graph is None:
            transition_matrix = np.stack(
                [multivariate_normal(
                    mean=center, cov=self.movement_var).pdf(
                    place_bin_centers)
                 for center in self.environment.place_bin_centers_], axis=1)
        else:
            place_bin_center_ind_to_node = np.asarray(
                self.environment.place_bin_centers_nodes_df_.node_id)
            transition_matrix = _random_walk_on_track_graph(
                self.environment.place_bin_centers_,
                self.movement_var,
                place_bin_center_ind_to_node,
                self.environment.distance_between_nodes_
            )

        is_track_interior = self.environment.is_track_interior_.ravel(order='F')
        transition_matrix[~is_track_interior] = 0.0
        transition_matrix[:, ~is_track_interior] = 0.0

        self.state_transition_ = _normalize_row_probability(transition_matrix)

        return self.state_transition_

@dataclass
class Uniform:
    environment_name: str
    environment2_name: str = None

    def make_state_transition(self, environments):
        self.environment1 = environments[self.environment_name]
        n_bins1 = self.environment1.place_bin_centers_.shape[0]
        is_track_interior1 = self.environment1.is_track_interior_.ravel(order='F')
        
        if self.environment2_name is None:
            n_bins2 = n_bins1
            is_track_interior2 = is_track_interior1.copy()
        else:
            self.environment2 = environments[self.environment2_name]
            n_bins2 = self.environment2.place_bin_centers_.shape[0]
            is_track_interior2 = self.environment2.is_track_interior_.ravel(order='F')

        transition_matrix = np.ones((n_bins1, n_bins2))

        transition_matrix[~is_track_interior1] = 0.0
        transition_matrix[:, ~is_track_interior2] = 0.0

        return _normalize_row_probability(transition_matrix)

@dataclass
class Identity:
    environment_name: str = None
        
    def make_state_transition(self, environments: dict):
        self.environment = environments[self.environment_name]
        n_bins = self.environment.place_bin_centers_.shape[0]
        
        transition_matrix = np.identity(n_bins)

        is_track_interior = self.environment.is_track_interior_.ravel(order='F')
        transition_matrix[~is_track_interior] = 0.0
        transition_matrix[:, ~is_track_interior] = 0.0

        return _normalize_row_probability(transition_matrix)

@dataclass
class EmpiricalMovement:
    environment_name: str = None
    speedup: int = 20

    def make_state_transition(self, environments, position, is_training=None):
        self.environment = environments[self.environment_name]
        
        if is_training is None:
            is_training = np.ones((position.shape[0]), dtype=np.bool)

        position = atleast_2d(position)[is_training]
        state_transition, _ = np.histogramdd(
            np.concatenate((position[1:], position[:-1]), axis=1),
            bins=self.environment.edges_ * 2, range=self.environment.position_range)
        original_shape = state_transition.shape
        n_position_dims = position.shape[1]
        shape_2d = np.product(original_shape[:n_position_dims])
        state_transition = _normalize_row_probability(
            state_transition.reshape((shape_2d, shape_2d), order='F'))
        return np.linalg.matrix_power(state_transition, self.speedup)
    

@dataclass
class RandomWalkDirection1():
    environment_name: str = None
    movement_var: float = 6.0

    def make_state_transition(self, environments: dict):
        random = (RandomWalk(self.environment_name, self.movement_var)
                  .make_state_transition(environments))
        return _normalize_row_probability(np.triu(random))


@dataclass
class RandomWalkDirection2():
    environment_name: str = None
    movement_var: float = 6.0

    def make_state_transition(self, environments: dict):
        random = (RandomWalk(self.environment_name, self.movement_var)
                  .make_state_transition(environments))
        return _normalize_row_probability(np.tril(random))

environments = {'w_track1': env1, 'w_track2': env2}
RandomWalk(environment_name='w_track2').make_state_transition(environments)
Uniform(environment_name='w_track1').make_state_transition(environments)
Identity(environment_name='w_track2').make_state_transition(environments)
EmpiricalMovement(environment_name='w_track2').make_state_transition(environments, position_df2.linear_position)

RandomWalkDirection1(environment_name='w_track2').make_state_transition(environments)
RandomWalkDirection1().make_state_transition(environments)
None == None
[Identity(environment_name='w_track2'), Identity(environment_name='w_track2'), RandomWalk(environment_name='w_track2')]
from dataclasses import dataclass

import numpy as np


@dataclass
class DiagonalDiscrete:
    diagonal_value: float = 0.98

    def make_state_transition(self, n_states):
        strong_diagonal = np.identity(n_states) * self.diagonal_value
        is_off_diag = ~np.identity(n_states, dtype=bool)
        strong_diagonal[is_off_diag] = (
            (1 - self.diagonal_value) / (n_states - 1))

        self.state_transition_ = strong_diagonal
        return self.state_transition_


@dataclass
class UniformDiscrete:
    def make_state_transition(self, n_states):
        self.state_transition_ = np.ones((n_states, n_states)) / n_states

        return self.state_transition_


@dataclass
class RandomDiscrete:
    def make_state_transition(self, n_states):
        state_transition = np.random.rand(n_states, n_states)
        state_transition /= state_transition.sum(axis=1, keepdims=True)

        self.state_transition_ = state_transition
        return self.state_transition_


UniformDiscrete().make_state_transition(3), DiagonalDiscrete(0.4).make_state_transition(3),  RandomDiscrete().make_state_transition(3)
DiagonalDiscrete(1/5).make_state_transition(5)
from replay_trajectory_classification.continuous_state_transitions import RandomWalk, Uniform
from replay_trajectory_classification.discrete_state_transitions import DiagonalDiscrete

continuous_state_transition_types = [[RandomWalk('w_track1', 6.0), Uniform('w_track1', 'w_track2')],
                                     [Uniform('w_track2', 'w_track1'), RandomWalk('w_track2', 6.0)]]

environments = {'w_track1': env1, 'w_track2': env2}
continuous_state_transition = [[ct.make_state_transition(environments) for ct in ct1]
                                for ct1 in continuous_state_transition_types]

discrete_state_transition = DiagonalDiscrete(0.98).make_state_transition(n_states=2)
discrete_state_transition
@dataclass
class UniformInitialConditions:
    environment_names: tuple
        
    def make_initial_conditions(self, environments):
        n_total_place_bins = 0
        initial_conditions = []
        
        for environment_name in self.environment_names:
            environment = environments[environment_name]
            n_place_bins = environment.is_track_interior_.sum()
            n_total_place_bins += n_place_bins
            initial_conditions.append(environment.is_track_interior_)
        
        return [ic / n_total_place_bins for ic in initial_conditions]

initial_conditions = UniformInitialConditions(('w_track1', 'w_track2')).make_initial_conditions(environments)
initial_conditions
from replay_trajectory_classification.spiking_likelihood import estimate_place_fields

environment_names_to_state = ('w_track1', 'w_track1', 'w_track1', 'w_track1')
encoding_groups_to_state =   ('Inbound', 'Inbound', 'Outbound', 'Outbound')

states = tuple(zip(environment_names_to_state, encoding_groups_to_state))

n_time = position_df1.linear_position.values.shape[0]
is_training1 = np.ones((n_time,), dtype=bool)
encoding_group_labels1 = np.concatenate((np.asarray(['Outbound'] * (n_time//2)), np.asarray(['Inbound'] * (n_time - (n_time//2)))))

n_time = position_df2.linear_position.values.shape[0]
is_training2 = np.ones((n_time,), dtype=bool)
encoding_group_labels2 = np.asarray(['Inbound'] * n_time)

training_data = {
    'w_track1': (position_df1.linear_position.values[:, np.newaxis], spikes1, is_training1, encoding_group_labels1),
    'w_track2': (position_df2.linear_position.values[:, np.newaxis], spikes2, is_training2, encoding_group_labels2),
}

place_fields = {}
for environment_name, encoding_group in set(states):
    print(environment_name, encoding_group)
    environment = environments[environment_name]
    position, spikes, is_training, encoding_group_labels = training_data[environment_name]
    
    place_fields[(environment_name, encoding_group)] = estimate_place_fields(
        position=position[is_training & (encoding_group_labels == encoding_group)],
        spikes=spikes[is_training & (encoding_group_labels == encoding_group)],
        place_bin_centers=environment.place_bin_centers_,
        place_bin_edges=environment.place_bin_edges_,
        penalty=1E1,
        knot_spacing=10)
    
place_fields
from replay_trajectory_classification.spiking_likelihood import estimate_place_fields

environment_names_to_state = ('w_track1', 'w_track2', 'w_track1', 'w_track2')
encoding_groups_to_state =   ('Inbound', 'Inbound', 'Outbound', 'Outbound')

states = tuple(zip(environment_names_to_state, encoding_groups_to_state))

n_time = position_df1.linear_position.values.shape[0]
is_training1 = np.ones((n_time,), dtype=bool)
encoding_group_labels1 = np.concatenate((np.asarray(['Outbound'] * (n_time//2)), np.asarray(['Inbound'] * (n_time - (n_time//2)))))
environment_labels1 = np.asarray(['w_track1'] * n_time)

n_time = position_df2.linear_position.values.shape[0]
is_training2 = np.ones((n_time,), dtype=bool)
encoding_group_labels2 = np.concatenate((np.asarray(['Outbound'] * (n_time//2)), np.asarray(['Inbound'] * (n_time - (n_time//2)))))
environment_labels2 = np.asarray(['w_track2'] * n_time)

position = np.concatenate((position_df1.linear_position.values[:, np.newaxis],
                           (position_df2.linear_position.values[:, np.newaxis])))
spikes = np.concatenate((spikes1, spikes2))
is_training = np.concatenate((is_training1, is_training2))
encoding_group_labels = np.concatenate((encoding_group_labels1, encoding_group_labels2))
environment_labels = np.concatenate((environment_labels1, environment_labels2))


place_fields = {}
for environment_name, encoding_group in set(states):
    print(environment_name, encoding_group)
    environment = environments[environment_name]
    
    is_encoding = (encoding_group_labels == encoding_group)
    is_environment = (environment_labels == environment_name)
    place_fields[(environment_name, encoding_group)] = estimate_place_fields(
        position=position[is_training & is_encoding & is_environment],
        spikes=spikes[is_training & is_encoding & is_environment],
        place_bin_centers=environment.place_bin_centers_,
        place_bin_edges=environment.place_bin_edges_,
        penalty=1E1,
        knot_spacing=10)
    
place_fields
@dataclass
class TrainingData:
    environment_name: str = None
    position: np.ndarray = None
    spikes: np.ndarray = None
    multiunit: np.ndarray = None
    is_training: np.ndarray = None
    encoding_group_labels: np.ndarray = None
        

@dataclass
class PredictData:
    position: np.ndarray = None
    spikes: np.ndarray = None
    multiunit: np.ndarray = None

        
training_data1 = TrainingData(
    environment_name='w_track1',
    position=position_df1.linear_position.values[:, np.newaxis],
    spikes=spikes1,
    is_training=is_training1,
    encoding_group_labels=encoding_group_labels1
)

training_data2 = TrainingData(
    environment_name='w_track2',
    position=position_df2.linear_position.values[:, np.newaxis],
    spikes=spikes2,
    is_training=is_training2,
    encoding_group_labels=encoding_group_labels2
)

# or just one long vector of environment, labels, position, spikes
fig, axes = plt.subplots(2, 1, sharex=True, sharey=True)
place_fields['w_track1', 'Inbound'].plot(x='position', hue='neuron', add_legend=False, ax=axes[0]);
place_fields['w_track1', 'Outbound'].plot(x='position', hue='neuron', add_legend=False, ax=axes[1]);
encoding_group_labels
def fit_place_fields(self, position, spikes, is_training=None,
                     encoding_group_labels=None,
                     encoding_group_to_state=None):
    logger.info('Fitting place fields...')
    if is_training is None:
        n_time = position.shape[0]
        is_training = np.ones((n_time,), dtype=np.bool)

    if encoding_group_labels is None:
        n_time = position.shape[0]
        encoding_group_labels = np.zeros((n_time,), dtype=np.int32)

    if encoding_group_to_state is None:
        n_states = len(self.continuous_transition_types)
        self.encoding_group_to_state_ = np.zeros(
            (n_states,), dtype=np.int32)
    else:
        self.encoding_group_to_state_ = np.asarray(encoding_group_to_state)

    is_training = np.asarray(is_training).squeeze()
    self.place_fields_ = []
    unique_labels = np.unique(encoding_group_labels[is_training])
    for encoding_group in unique_labels:
        self.place_fields_.append(estimate_place_fields(
            position=position[is_training & (
                encoding_group_labels == encoding_group)],
            spikes=spikes[is_training & (
                encoding_group_labels == encoding_group)],
            place_bin_centers=self.place_bin_centers_,
            place_bin_edges=self.place_bin_edges_,
            penalty=self.spike_model_penalty,
            knot_spacing=self.knot_spacing))
# how specify training data
# how to incorporate local states
# adding information by multiplying likelihoods (LFPs)
argh = [[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]]
blah = []

for row in argh:
    blah.append([])
    for column in row:
        blah[-1].append(column)
blah<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09618112097494354, 'fraction_numerical': 0.013605084834795398, 'mean_word_length': 3.2562038900067067, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '32679', 'license': 'MIT', 'n_tokens_mistral': 5968, 'n_tokens_neox': 5763, 'n_words': 953, 'path': 'notebooks/work_in_progress/Two_Environments.ipynb', 'repo_name': 'Eden-Kramer-Lab/replay_trajectory_classification', 'size': '19037'}