# Semi-supervized VAE StarNet

* Inputs: APOGEE DR12 real spectra and simulated spectra with ASSET
* Outputs: stellar parameters from ASSET

This notebook goes through the building and training of two VAE models in series.

```python
import numpy as np
import time
import h5py
import keras
import matplotlib.pyplot as plt
import math

import sys

from keras.layers import (Input, Dense, Lambda, Flatten, Reshape, BatchNormalization, Activation,
                          Dropout, Conv1D, UpSampling1D, MaxPooling1D, ZeroPadding1D, LeakyReLU, AveragePooling1D)
from keras.layers.merge import concatenate
from keras.optimizers import Adam
from keras.models import Model
from keras import metrics
from keras import backend as K

from keras.engine.topology import Layer

from keras import regularizers

import keras.initializers
from keras_contrib.layers import InstanceNormalization

plt.switch_backend('agg')
```

### Data loading and normalization

Some helper functions for the data.

```python
def load_train_data_weighted_with_asset(data_file_l,data_file_u,indices_l=None,indices_u=None,
                                       unlabeled=True,labeled=True):
    # load APOGEE spectra
    if unlabeled:
        # Grab all
        if indices_u is None:
            with h5py.File(data_file_u,"r") as F:
                ap_spectra = F['APOGEE spectrum'][:]
                ap_err_spectra = F['APOGEE error_spectrum'][:]
                ap_labels = np.column_stack((F['APOGEE TEFF'][:],F['APOGEE LOGG'][:],F['APOGEE FE_H'][:],
                                            F['APOGEE alpha_M'][:],F['APOGEE N_M'][:],F['APOGEE C_M'][:]))
        # Grab a batch
        else:
            with h5py.File(data_file_u, "r") as F:
                indices_bool = np.ones((len(F['APOGEE spectrum']),),dtype=bool)
                indices_bool[:] = False
                indices_bool[indices_u] = True

                ap_spectra = F['APOGEE spectrum'][indices_bool,:]
                ap_err_spectra = F['APOGEE error_spectrum'][indices_bool,:]
                ap_labels = np.column_stack((F['APOGEE TEFF'][indices_bool],F['APOGEE LOGG'][indices_bool],
                                         F['APOGEE FE_H'][indices_bool], F['APOGEE alpha_M'][indices_bool],
                                         F['APOGEE N_M'][indices_bool],F['APOGEE C_M'][indices_bool]))
    else:
        ap_spectra=[]
        ap_err_spectra=[]
        ap_labels=[]

    # load ASSET spectra
    if labeled:
        # Grab all
        if indices_l is None:
            with h5py.File(data_file_l,"r") as F:
                synth_spectra = F['spectrum'][:]
                synth_labels = np.column_stack((F['teff'][:],F['logg'][:],F['M_H'][:],
                                                F['a_M'][:],F['N_M'][:],F['C_M'][:]))

                bad_spectra = list(set(list(np.where(synth_spectra<0.)[0])))
                for index in sorted(bad_spectra, reverse=True):
                    del synth_spectra[index]
                    del synth_labels[index]
        # Grab a batch
        else:
            with h5py.File(data_file_l, "r") as F:
                indices_bool = np.ones((len(F['spectrum']),),dtype=bool)
                indices_bool[:] = False
                indices_bool[indices_l] = True

                synth_spectra = F['spectrum'][indices_bool,:]
                synth_labels = np.column_stack((F['teff'][indices_bool],F['logg'][indices_bool],F['M_H'][indices_bool],
                                                F['a_M'][indices_bool],F['N_M'][indices_bool],F['C_M'][indices_bool]))

                synth_spectra=list(synth_spectra)
                synth_labels=list(synth_labels)
                bad_spectra = list(set(list(np.where(synth_spectra<0.)[0])))
                for index in sorted(bad_spectra, reverse=True):
                    del synth_spectra[index]
                    del synth_labels[index]
                synth_spectra=np.array(synth_spectra)
                synth_labels=np.array(synth_labels)
    else:
        synth_spectra=[]
        synth_labels=[]

    return ap_spectra,ap_err_spectra,ap_labels,synth_spectra,synth_labels

# normalize label variables to have zero-mean and unit-variance
def normalize(lb,mean_and_std_file='/data/stars/apogee/dr12/asset_labels_mean_and_std.npy'):
    mean_and_std = np.load(mean_and_std_file)
    return (lb-mean_and_std[0])/mean_and_std[1]

def denormalize(lb_norm,mean_and_std_file='/data/stars/apogee/dr12/asset_labels_mean_and_std.npy'):
    mean_and_std = np.load(mean_and_std_file)
    return ((lb_norm*mean_and_std[1]) + mean_and_std[0])

# normalize each spectrum to have zero-mean and unit-variance (ignore zero-values whe calculating mean and std)
def normalize_spec(specs):
    specs[specs == 0] = np.nan
    mean_spec = np.nanmean(specs,axis=1).reshape(len(specs),1)
    std_spec = np.nanstd(specs,axis=1).reshape(len(specs),1)
    specs[np.isnan(specs)] = 0.
    return (specs-mean_spec)/std_spec
```

```python
# function for reshaping spectra into appropriate format for CNN
def cnn_reshape(spectra):
    return spectra.reshape(spectra.shape[0],spectra.shape[1],1)
```

```python
# Set some model hyper-parameters

img_cols, img_chns = 7214, 1
num_fluxes=7214
input_shape=(num_fluxes,1)
z_dims = 16
num_labels=6
batch_size = 64
epsilon_std = 1.0
learning_rate = 0.001
decay = 0.0
var_epsilon = 0.010

padding=u'same'
kernel_init = 'he_normal'
bias_init = keras.initializers.Zeros()

num_filters = [16, 16]
filter_length = 3
pool_length = 4
num_hidden = [256,128]
```

## Zero-Augmentation

In an effort to evenly distribute the weighting of the VAE, throughout training, a *zero-augmentation* technique was applied to the training spectra samples - both synthetic and observed.

The zero-augmentation is implemented as the first layer in the encoder where a zero-augmentation mask is sent as an input along with the input spectrum and the two are multiplied together.

The zero-augmentation mask is the same size as the input spectrum vector and is composed of ones and zeros. For the APOGEE wave-grid, the spectral region is divided into seven *chunks* and for each input spectrum a random 0-3 of these chunks are assigned to be zeros while the remainder of the zero-augmentation mask is made up of ones. This means for a given spectrum, the input for training may be 4/7ths, 5/7ths, 6/7ths, or the entire spectrum. This augmentation is done randomly throughout training, meaning that each spectrum will be randomly assigned a different zero-augmentation mask at every iteration.

```python
# zero-augmentation layer
class ZeroAugmentLayer(Layer):
    def __init__(self, **kwargs):
        self.is_placeholder = True
        super(ZeroAugmentLayer, self).__init__(**kwargs)

    def zero_agument(self, x_real, zero_mask):
        return x_real*zero_mask

    def call(self, inputs):
        x_real = inputs[0]
        zero_mask = inputs[1]
        x_augmented = self.zero_agument(x_real, zero_mask)

        return x_augmented

# a function for creating the zero-masks used during training
def create_zero_mask(spectra,min_chunks,max_chunks,chunk_size,dataset=None,ones_padded=False):
    if dataset is None:
        zero_mask = np.ones_like(spectra)
    elif dataset=='apogee':
        zero_mask = np.ones((spectra.shape[0],7214))
    elif dataset=='segue':
        zero_mask = np.ones((spectra.shape[0],3688))

    num_spec = zero_mask.shape[0]
    len_spec = zero_mask.shape[1]
    num_bins = len_spec/chunk_size
    remainder = len_spec%chunk_size
    spec_sizes = np.array([chunk_size for i in range(num_bins)])
    spec_sizes[-1]=spec_sizes[-1]+remainder

    num_bins_removed = np.random.randint(min_chunks,max_chunks+1,size=(num_spec,))
    for i, mask in enumerate(zero_mask):
        bin_indx_removed = np.random.choice(num_bins, num_bins_removed[i], replace=False)
        for indx in bin_indx_removed:
            if indx==0:
                mask[indx*spec_sizes[indx]:(indx+1)*spec_sizes[indx]]=0.
            else:
                mask[indx*spec_sizes[indx-1]:indx*spec_sizes[indx-1]+spec_sizes[indx]]=0.

    return zero_mask
```

```python
def set_trainable(model, train):
    """
    Enable or disable training for the model
    """
    model.trainable = train
    for l in model.layers:
        l.trainable = train
```

# Build encoder

Takes spectra (x) and zero-augmentation mask as inputs and outputs latent distribution (z_mean and z_log_var). If the encoder is for the labeled data, the outputs will also include the stellar parameter distribution (y_mean and y_log_var)

```python
def build_encoder(input_spec,input_mask,labeled=False):

    # input conv layer with filter length 1, no bias value
    #x = Conv1D(kernel_initializer=keras.initializers.Constant(0.5), activation='linear', padding="same", filters=1,
           #kernel_size=1,use_bias=False)(input_spec)
    # instance normalize to bring each spectrum to zero-mean and unit variance
    #spec_normed = InstanceNormalization()(input_spec)

    x = ZeroAugmentLayer()([input_spec,input_mask])

    # first conv layer
    x = Conv1D(filters=num_filters[0], kernel_size=filter_length, strides=1, kernel_initializer=kernel_init,
                   bias_initializer=bias_init, padding=padding)(x)
    x = Activation('relu')(x)
    x = InstanceNormalization()(x)

    # second conv layer
    x = Conv1D(filters=num_filters[1], kernel_size=filter_length, strides=1, kernel_initializer=kernel_init,
                   bias_initializer=bias_init, padding=padding)(x)
    x = Activation('relu')(x)
    x = InstanceNormalization()(x)

    # Avg pooling
    x = AveragePooling1D(pool_size=pool_length)(x)
    x = Flatten()(x)
    x = Dropout(0.2)(x)

    # intermediate dense block
    x = Dense(num_hidden[0], kernel_initializer=kernel_init)(x)
    x = Activation('relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(0.3)(x)

    # intermediate dense block
    x = Dense(num_hidden[1], kernel_initializer=kernel_init)(x)
    x = Activation('relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(0.2)(x)

    # latent distribution output
    z_mean = Dense(z_dims)(x)
    z_log_var = Dense(z_dims)(x)

    if labeled:
        # latent distribution output
        y_mean = Dense(num_labels,name='y_mean')(x)
        y_log_var = Dense(num_labels,name='y_log_var')(x)

        # return model as well as normalized spectra tensor
        return Model([input_spec,input_mask],[z_mean,z_log_var,y_mean,y_log_var]),input_spec
    else:
        # return model as well as normalized spectra tensor
        return Model([input_spec,input_mask],[z_mean,z_log_var]),input_spec
```

```python
# function for obtaining a latent sample given a distribution
def sampling(args, epsilon_std=epsilon_std):
    z_mean, z_log_var = args

    epsilon = K.random_normal(shape=(K.int_shape(z_mean)[1],),
                              mean=0., stddev=epsilon_std)
    return z_mean + K.exp(z_log_var) * epsilon
```

# Build decoder

- if labeled input model: takes  z (latent variables) and y (stellar labels) concatenated into a single vector as an input and outputs a *normalized* stellar spectrum
- if unlabeled input model: takes  z (latent variables) as an input and outputs a *normalized* stellar spectrum

```python
def build_decoder(input_zy):

    # input fully-connected block
    x = Dense(num_hidden[1])(input_zy)
    x = Activation('relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(0.2)(x)

    # upsampling fully-connected block
    x = Dense(input_shape[0] * num_filters[1])(x)
    x = Activation('relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(0.2)(x)

    # reshape for convolutional blocks
    output_shape = (input_shape[0], num_filters[1],)
    x = Reshape(output_shape)(x)

    # first deconv block
    x = Conv1D(kernel_initializer=kernel_init,bias_initializer=bias_init,padding="same",
                   filters=num_filters[1],kernel_size=filter_length)(x)
    x = Activation('relu')(x)
    x = InstanceNormalization()(x)
    x = Dropout(0.1)(x)

    # second deconv block
    x = Conv1D(kernel_initializer=kernel_init,bias_initializer=bias_init,padding="same",
                   filters=num_filters[0],kernel_size=filter_length)(x)
    x = Activation('relu')(x)
    x = InstanceNormalization()(x)

    # output conv layer
    x = Conv1D(kernel_initializer=kernel_init,bias_initializer=bias_init,padding="same",
                   filters=1,kernel_size=filter_length,activation='linear')(x)

    return Model(input_zy,x)
```

# Build VAE_A

This model takes spectra (spec_A) as inputs, encodes (encoder_A) them into latent variables (z_A) and stellar labels (y_A) and then concatenates z_A and y_A into vector input for decoder_A that produces a spectrum.

```python
# encoder_A input placeholders
input_spec_A = Input(shape=input_shape)
input_mask_A = Input(shape=input_shape)

# error_spectra_A placeholder
input_err_spec_A = Input(shape=input_shape)

# predictor_input_A placeholder
input_z_A = Input(shape=(z_dims,))

# decoder_input_A placeholder
input_zy_A = Input(shape=(z_dims+num_labels,))

# encoder model A
encoder_A, normed_spec_A = build_encoder(input_spec_A, input_mask_A,labeled=True)

# decoder model A
decoder_A = build_decoder(input_zy_A)
```

# Build VAE_B

This model takes spectra (spec_B) as inputs, encodes (encoder_B) them into latent variables (z_B), uses those latent variables as the input for decoder_B that produces a spectrum.

```python
# encoder_input_B placeholders
input_spec_B = Input(shape=input_shape)
input_mask_B = Input(shape=input_shape)

# error_spectra_B placeholder
input_err_spec_B = Input(shape=input_shape)

# decoder_input_B placeholder
input_z_B = Input(shape=(z_dims,))

# encoder model B
encoder_B, normed_spec_B = build_encoder(input_spec_B, input_mask_B)

# decoder model B
decoder_B = build_decoder(input_z_B)
```

## Create loss functions

The VAEs have two loss functions that are minimized simultaneously:

1. a weighted mean-squared-error to analyze the predicted spectra:
\begin{equation}
mse = \frac{1}{N}\sum{\frac{(x_{true}-x_{pred})^2}{(x_{err})^2}}
\end{equation}

2. a relative entropy, KL (Kullback–Leibler divergence) loss to keep the latent variables within a similar distribuition:
\begin{equation}
KL = \frac{1}{N}\sum{-\frac{1}{2}(1.0+z_{log\_var} - z_{avg}^2 - e^{z_{log\_var}})}
\end{equation}

```python
class VAE_Labeled_LossLayer_weighted(Layer):
    __name__ = u'vae_labeled_loss_layer'

    def __init__(self, **kwargs):
        self.is_placeholder = True
        super(VAE_Labeled_LossLayer_weighted, self).__init__(**kwargs)

    def lossfun(self, x_true, x_pred, z_avg_1, z_log_var_1, z_avg_2, z_log_var_2,x_err,y_avg,y_log_var):
        rec_loss = K.mean(K.square((x_true - x_pred)/x_err))
        kl_loss_1 = K.mean(-0.5 * K.sum(1.0 + z_log_var_1 - K.square(z_avg_1) - K.exp(z_log_var_1), axis=-1))
        kl_loss_2 = K.mean(-0.5 * K.sum(1.0 + z_log_var_2 - K.square(z_avg_2) - K.exp(z_log_var_2), axis=-1))
        kl_loss_3 = K.mean(-0.5 * K.sum(1.0 + y_log_var - K.square(y_avg) - K.exp(y_log_var), axis=-1))
        return rec_loss + kl_loss_1 +kl_loss_2
        # not sure which one to use
        #return rec_loss + kl_loss_1 +kl_loss_2+kl_loss_3

    def call(self, inputs):
        x_true = inputs[0]
        x_pred = inputs[1]
        x_err = inputs[2]
        z_avg_1 = inputs[3]
        z_log_var_1 = inputs[4]
        z_avg_2 = inputs[5]
        z_log_var_2 = inputs[6]
        y_avg = inputs[7]
        y_log_var = inputs[8]
        loss = self.lossfun(x_true, x_pred, z_avg_1, z_log_var_1, z_avg_2, z_log_var_2, x_err, y_avg, y_log_var)
        self.add_loss(loss, inputs=inputs)

        return x_true

class VAE_Unlabeled_LossLayer_weighted(Layer):
    __name__ = u'vae_labeled_loss_layer'

    def __init__(self, **kwargs):
        self.is_placeholder = True
        super(VAE_Unlabeled_LossLayer_weighted, self).__init__(**kwargs)

    def lossfun(self, x_true, x_pred, z_avg_1, z_log_var_1, z_avg_2, z_log_var_2,x_err):
        rec_loss = K.mean(K.square((x_true - x_pred)/x_err))
        kl_loss_1 = K.mean(-0.5 * K.sum(1.0 + z_log_var_1 - K.square(z_avg_1) - K.exp(z_log_var_1), axis=-1))
        kl_loss_2 = K.mean(-0.5 * K.sum(1.0 + z_log_var_2 - K.square(z_avg_2) - K.exp(z_log_var_2), axis=-1))
        return rec_loss + kl_loss_1 +kl_loss_2

    def call(self, inputs):
        x_true = inputs[0]
        x_pred = inputs[1]
        x_err = inputs[2]
        z_avg_1 = inputs[3]
        z_log_var_1 = inputs[4]
        z_avg_2 = inputs[5]
        z_log_var_2 = inputs[6]
        loss = self.lossfun(x_true, x_pred, z_avg_1, z_log_var_1, z_avg_2, z_log_var_2, x_err)
        self.add_loss(loss, inputs=inputs)

        return x_true

# dummy loss to give zeros, hence no gradients to train
# the real loss is computed as the layer shown above and therefore this dummy loss is just
# used to satisfy keras notation when compiling the model
def zero_loss(y_true, y_pred):
    return K.zeros_like(y_true)
```

# Labeled process

- input_spec_A: labeled spectra
- input_mask_A: labeled zero-mask
- input_y_l: known stellar labels

Process:
    1. encode input_spec_A using encoder_A to produce latent distribution (z_mean_A_l, z_log_var_A_l) and stellar parameter distribution (y_mean_A_l, y_log_var_A_l)
    2. sample from latent distribution to create a vector (z_A_l) of length z_dims
    3. same for stellar parameter distribution
    4. merge latent samples with stellar labels to create a vector (merged_zy_l) of length z_dims+num_labels
    5. decode merged_zy_l using decoder_A into spectrum (output_spec_A_l)
    6. encode output_spec_A_l using encoder_B to produce latent distribution (z_mean_B_l, z_log_var_B_l)
    7. sample from latent distribution to create a vector (z_B_l) of length z_dims
    8. decode z_B_l using decoder_B into spectrum (output_spec_B_l)

```python
# outputs for encoder_A
z_mean_A_l, z_log_var_A_l, y_mean_A_l, y_log_var_A_l = encoder_A([input_spec_A, input_mask_A])

# loss for predictor
def logy_loss(y, y_pred,y_log_var=y_log_var_A_l):
    loss = (  0.5 * math.log(2 * math.pi)
            + 0.5 * K.log(K.exp(y_log_var) + var_epsilon)
            + 0.5 * K.square(y - y_pred) / (K.exp(y_log_var) + var_epsilon))
    loss = num_labels*K.sum(loss, axis=-1)
    return K.mean(loss)

# sample from latent distribution given z_mean_A_l and z_log_var_A_l
z_A_l = Lambda(sampling, output_shape=(z_dims,))([z_mean_A_l, z_log_var_A_l])
y_A_l = Lambda(sampling, output_shape=(num_labels,))([y_mean_A_l, y_log_var_A_l])

# labeled known stellar parameters input placeholder
input_y_l = Input(shape=(num_labels,))

# merge latent variables with known stellar labels
merged_zy_l = concatenate([y_A_l, z_A_l])

# decode latent variables and known stellar labels, input_y_l, into stellar spectrum
output_spec_A_l = decoder_A(merged_zy_l)

# encode produced spectrum output_spec_A_l using encoder B
z_mean_B_l, z_log_var_B_l = encoder_B([output_spec_A_l, input_mask_A])

# sample from latent distribution given z_mean_B_l and z_log_var_B_l
z_B_l = Lambda(sampling, output_shape=(z_dims,))([z_mean_B_l, z_log_var_B_l])

# decode latent variables, z_B_l, into stellar spectrum
output_spec_B_l = decoder_B(z_B_l)
```

```python
# create loss layer that will evaluate:
# 1. the normalized spectrum from the encoder_A input layer against
#         the output normalized spectrum produced by decoder_B
# 2. The two latent distributions using the KL loss
vae_l_loss = VAE_Labeled_LossLayer_weighted()([normed_spec_A, output_spec_B_l, input_err_spec_A, z_mean_A_l, z_log_var_A_l,
                                               z_mean_B_l, z_log_var_B_l, y_mean_A_l, y_log_var_A_l])

# create model that takes labeled spectra, error spectra, zero-augmentation masks, and known stellar labels as inputs
#      and outputs the above loss layer, as well as the predicted stellar parameters
labeled_vae = Model(inputs=[input_spec_A, input_mask_A, input_err_spec_A, input_y_l], outputs=[vae_l_loss, y_mean_A_l])

# compile model using zero-loss for the first output (loss already computed in loss layer)
#      and a mean-squared-error to evaluate the predicted stellar labels
labeled_vae.compile(loss=[zero_loss,logy_loss],
                         optimizer=Adam(lr=1.0e-4, beta_1=0.5))

labeled_vae.summary()
```

    __________________________________________________________________________________________________
    Layer (type)                    Output Shape         Param #     Connected to
    ==================================================================================================
    input_1 (InputLayer)            (None, 7214, 1)      0
    __________________________________________________________________________________________________
    input_2 (InputLayer)            (None, 7214, 1)      0
    __________________________________________________________________________________________________
    model_1 (Model)                 [(None, 16), (None,  7426304     input_1[0][0]
                                                                     input_2[0][0]
    __________________________________________________________________________________________________
    lambda_2 (Lambda)               (None, 6)            0           model_1[1][2]
                                                                     model_1[1][3]
    __________________________________________________________________________________________________
    lambda_1 (Lambda)               (None, 16)           0           model_1[1][0]
                                                                     model_1[1][1]
    __________________________________________________________________________________________________
    concatenate_1 (Concatenate)     (None, 22)           0           lambda_2[0][0]
                                                                     lambda_1[0][0]
    __________________________________________________________________________________________________
    model_2 (Model)                 (None, 7214, 1)      15356469    concatenate_1[0][0]
    __________________________________________________________________________________________________
    model_3 (Model)                 [(None, 16), (None,  7424756     model_2[1][0]
                                                                     input_2[0][0]
    __________________________________________________________________________________________________
    lambda_3 (Lambda)               (None, 16)           0           model_3[1][0]
                                                                     model_3[1][1]
    __________________________________________________________________________________________________
    model_4 (Model)                 (None, 7214, 1)      15355701    lambda_3[0][0]
    __________________________________________________________________________________________________
    input_3 (InputLayer)            (None, 7214, 1)      0
    __________________________________________________________________________________________________
    vae__labeled__loss_layer_weight [(None, 7214, 1), (N 0           input_1[0][0]
                                                                     model_4[1][0]
                                                                     input_3[0][0]
                                                                     model_1[1][0]
                                                                     model_1[1][1]
                                                                     model_3[1][0]
                                                                     model_3[1][1]
                                                                     model_1[1][2]
                                                                     model_1[1][3]
    ==================================================================================================
    Total params: 45,563,230
    Trainable params: 45,099,486
    Non-trainable params: 463,744
    __________________________________________________________________________________________________

# Unlabeled process

First freeze y_mean and y_log_var model weights

- input_spec_B: unlabeled spectra
- input_mask_B: unlabeled zero-mask

Process:
    1. encode input_spec_B using encoder_B to produce latent distribution (z_mean_B_u, z_log_var_B_u)
    2. sample from latent distribution to create a vector (z_B_u) of length z_dims
    3. decode z_B_u using decoder_B into spectrum (output_spec_B_u)
    4. encode output_spec_B_u using encoder_A to produce latent distribution (z_mean_A_u, z_log_var_A_u) and stellar parameter distribution (y_mean_A_u, y_log_var_A_u)
    5. sample from latent distribution to create a vector (z_A_u) of length z_dims
    6. same for stellar parameter distribution
    7. merge latent samples with predicted stellar labels to create a vector (merged_zy_u) of length z_dims+num_labels
    8. decode merged_zy_u using decoder_A into spectrum (output_spec_A_u)

```python
encoder_A.get_layer('y_mean').trainable=False
encoder_A.get_layer('y_log_var').trainable=False
```

```python
# outputs for encoder_B
z_mean_B_u, z_log_var_B_u = encoder_B([input_spec_B, input_mask_B])

# sample from latent distribution given z_mean_B_u and z_log_var_B_u
z_B_u = Lambda(sampling, output_shape=(z_dims,))([z_mean_B_u, z_log_var_B_u])

# decode latent variables, z_B_u, into stellar spectrum
output_spec_B_u = decoder_B(z_B_u)

# encode output_spec_B_u with encoder_A
z_mean_A_u, z_log_var_A_u, y_mean_A_u, y_log_var_A_u = encoder_A([output_spec_B_u, input_mask_B])

# sample from latent distribution given z_mean_A_u and z_log_var_A_u
z_A_u = Lambda(sampling, output_shape=(z_dims,))([z_mean_A_u, z_log_var_A_u])
y_A_u = Lambda(sampling, output_shape=(z_dims,))([y_mean_A_u, y_log_var_A_u])

# merge latent variables with predicted stellar labels
merged_zy_u = concatenate([y_A_u, z_A_u])

# decode latent variables and known stellar labels, input_y_u, into stellar spectrum
output_spec_A_u = decoder_A(merged_zy_u)
```

```python
# create loss layer that will evaluate:
# 1. the normalized spectrum from the encoder_B input layer against
#         the output normalized spectrum produced by decoder_A
# 2. The two (or three) latent distributions using the KL loss
vae_u_loss = VAE_Unlabeled_LossLayer_weighted()([normed_spec_B, output_spec_A_u, input_err_spec_B, z_mean_A_u, z_log_var_A_u,
                                               z_mean_B_u, z_log_var_B_u])

# create model that takes unlabeled spectra, error spectra, and zero-augmentation masks as inputs
#      and outputs the above loss layer
unlabeled_vae = Model(inputs=[input_spec_B, input_mask_B, input_err_spec_B], outputs=[vae_u_loss])

# compile model using zero-loss for the first output (loss already computed in loss layer)
unlabeled_vae.compile(loss=[zero_loss],
                         optimizer=Adam(lr=1.0e-4, beta_1=0.5))

unlabeled_vae.summary()
```

    __________________________________________________________________________________________________
    Layer (type)                    Output Shape         Param #     Connected to
    ==================================================================================================
    input_6 (InputLayer)            (None, 7214, 1)      0
    __________________________________________________________________________________________________
    input_7 (InputLayer)            (None, 7214, 1)      0
    __________________________________________________________________________________________________
    model_3 (Model)                 [(None, 16), (None,  7424756     input_6[0][0]
                                                                     input_7[0][0]
    __________________________________________________________________________________________________
    lambda_4 (Lambda)               (None, 16)           0           model_3[2][0]
                                                                     model_3[2][1]
    __________________________________________________________________________________________________
    model_4 (Model)                 (None, 7214, 1)      15355701    lambda_4[0][0]
    __________________________________________________________________________________________________
    model_1 (Model)                 [(None, 16), (None,  7426304     model_4[2][0]
                                                                     input_7[0][0]
    __________________________________________________________________________________________________
    lambda_6 (Lambda)               (None, 16)           0           model_1[2][2]
                                                                     model_1[2][3]
    __________________________________________________________________________________________________
    lambda_5 (Lambda)               (None, 16)           0           model_1[2][0]
                                                                     model_1[2][1]
    __________________________________________________________________________________________________
    concatenate_2 (Concatenate)     (None, 32)           0           lambda_6[0][0]
                                                                     lambda_5[0][0]
    __________________________________________________________________________________________________
    model_2 (Model)                 (None, 7214, 1)      15356469    concatenate_2[0][0]
    __________________________________________________________________________________________________
    input_8 (InputLayer)            (None, 7214, 1)      0
    __________________________________________________________________________________________________
    vae__unlabeled__loss_layer_weig [(None, 7214, 1), (N 0           input_6[0][0]
                                                                     model_2[2][0]
                                                                     input_8[0][0]
                                                                     model_1[2][0]
                                                                     model_1[2][1]
                                                                     model_3[2][0]
                                                                     model_3[2][1]
    ==================================================================================================
    Total params: 45,563,230
    Trainable params: 45,097,938
    Non-trainable params: 465,292
    __________________________________________________________________________________________________

```python
# models that encode and then decode spectra (this is used to plot the intermediate results during training)
gen_x_to_x_l = Model([input_spec_A,input_mask_A,input_y_l], [normed_spec_A,output_spec_B_l])
gen_x_to_x_l.compile(loss=['mse','mse'],
                         optimizer=Adam(lr=1.0e-4, beta_1=0.5))

gen_x_to_x_u = Model([input_spec_B,input_mask_B], [normed_spec_B,output_spec_A_u])
gen_x_to_x_u.compile(loss=['mse','mse'],
                         optimizer=Adam(lr=1.0e-4, beta_1=0.5))
```

```python
# model that predicts stellar labels for unlabeled data
# (this is used to make sure the error is generally decreasing throughout trainig)
predictor_u = Model([input_spec_B,input_mask_B], y_mean_A_u)
```

```python
def time_format(t):
    m, s = divmod(t, 60)
    m = int(m)
    s = int(s)
    if m == 0:
        return u'%d sec' % s
    else:
        return u'%d min' % (m)

# function for training
def train_on_batch(x_obs,x_err_obs,x_synth,y_synth):

    # add noise to synthetic data
    noise_factor = 0.03
    x_err_synth = noise_factor * np.random.normal(loc=0.0, scale=1.0, size=x_synth.shape)
    x_err_synth[x_err_synth<0.005]=0.005
    x_synth = x_synth + x_err_synth

    # normalize each spectrum to have zero-mean and unit variance
    x_synth = normalize_spec(x_synth)
    x_obs = normalize_spec(x_obs)

    # Labeled training
    zero_mask = create_zero_mask(x_synth,
                                          0,3,1030,dataset=None,ones_padded=False)
    loss = labeled_vae.train_on_batch([cnn_reshape(x_synth),
                                     cnn_reshape(zero_mask), cnn_reshape(x_err_synth), y_synth],
                                    [cnn_reshape(x_synth), y_synth])

    # Unlabeled training
    zero_mask = create_zero_mask(x_obs,0,3,1030,dataset=None,ones_padded=False)
    loss += [unlabeled_vae.train_on_batch([cnn_reshape(x_obs),
                                           cnn_reshape(zero_mask),cnn_reshape(x_err_obs)],
                                          cnn_reshape(x_obs))]

    losses = {'l_loss': loss[0],
                 'l2_loss': loss[1],
                 'p_loss': loss[2],
                 'u_loss': loss[3],}
    return losses
```

## Main training loop

```python
def fit_model(model_name, data_file_l, data_file_u, epochs, reporter):

    # check to see how many spectra are in the unlabelled data file
    with h5py.File(data_file_u, "r") as F:
            num_data_ap = len(F['APOGEE spectrum'])

    # let's use 90% of the spectra
    num_data_train_ap = int(num_data_ap*0.9)
    test_indices_range_ap = [num_data_train_ap,num_data_ap]

    # check to see how many spectra are in the labelled data file
    with h5py.File(data_file_l, "r") as F:
            num_data_synth = len(F['spectrum'])

    # let's use 90% of the spectra
    num_data_train_synth = int(num_data_synth*0.9)
    test_indices_range_synth = [num_data_train_synth,num_data_synth]

    # load training spectra
    ap_spectra, ap_err_spectra, ap_labels, synth_spectra, synth_labels = load_train_data_weighted_with_asset(data_file_l, data_file_u, indices_l=np.arange(0,num_data_train_synth), indices_u=np.arange(0,num_data_train_ap))

    # grab good labels from unlabeled dataset to compare predictions to throughout training
    good_indices = np.where((ap_labels[:,0]!=-9999.)&(ap_labels[:,1]!=-9999.)&(ap_labels[:,2]!=-9999.)&
                            (ap_labels[:,3]!=-9999.)&(ap_labels[:,4]!=-9999.)&(ap_labels[:,5]!=-9999.))[0]

    num_data_train_ap = len(ap_spectra)
    perm_ap = np.random.permutation(num_data_train_ap)

    # loop through the number of epochs
    for e in xrange(start_e,epochs):

        perm_synth = np.random.permutation(num_data_train_synth)

        start_time = time.time()

        # loop through the batches
        losses_=[]
        for b in xrange(0, 200, batchsize): #num_data_train_synth

            # check current batchsize
            bsize = min(batchsize, num_data_train_synth - b)
            indx_synth = perm_synth[b:b+bsize]

            if len(perm_ap)<bsize:
                perm_ap = np.random.permutation(num_data_train_ap)
            indx_ap = perm_ap[0:bsize]
            perm_ap = perm_ap[bsize:]

            # grab a batch of unlabeled data
            x_obs = np.copy(ap_spectra[indx_ap])
            x_err_obs = np.copy(ap_err_spectra[indx_ap])

            # grab a batch of labeled data
            x_synth = np.copy(synth_spectra[indx_synth])
            y_synth = np.copy(synth_labels[indx_synth])

            # normalize stellar parameters
            y_synth = normalize(y_synth)

            # train on batch
            losses = train_on_batch(x_obs,x_err_obs,x_synth,y_synth)

            losses_.append(losses)
            # Print current status
            ratio = 100.0 * (b + bsize) / num_data_train_synth
            print unichr(27) + u"[2K",; sys.stdout.write(u'')
            print u'\rEpoch #%d | %d / %d (%6.2f %%) ' %                   (e + 1, b + bsize, num_data_train_synth, ratio),; sys.stdout.write(u'')

            for k in reporter:
                if k in losses:
                    print u'| %s = %5.3f ' % (k, losses[k]),; sys.stdout.write(u'')

            # Compute ETA
            elapsed_time = time.time() - start_time
            eta = elapsed_time / (b + bsize) * (num_data_train_synth - (b + bsize))
            print u'| ETA: %s ' % time_format(eta),; sys.stdout.write(u'')

            sys.stdout.flush()

        print u''

        # Print epoch status
        ratio = 100.0
        print unichr(27) + u"[2K",; sys.stdout.write(u'')
        print u'\rEpoch #%d | %d / %d (%6.2f %%) ' %               (e + 1, num_data_train_synth, num_data_train_synth, ratio),; sys.stdout.write(u'')

        losses_all = {}
        for k in losses_[0].iterkeys():
            losses_all[k] = tuple(d[k] for d in losses_)
        for k in reporter:
            if k in losses_all:
                losses_all[k]=np.sum(losses_all[k])/len(losses_)
        for k in reporter:
            if k in losses_all:
                print u'| %s = %5.3f ' % (k, losses_all[k]),; sys.stdout.write(u'')

        # save losses to file to analyze throughout training
        myfile = open(model_name+'.txt', 'a')
        for k in reporter:
            if k in losses_all:
                myfile.write("%s," % losses_all[k])
        myfile.write("\n")
        myfile.close()

        # compute mean-squared error on observed spec (compared against original pipeline)
        sample_indices = np.random.choice(good_indices, 1000, replace=False)
        samples = np.copy(ap_spectra[sample_indices])
        samples = normalize_spec(samples)
        mse = ((predictor_u.predict([cnn_reshape(samples),cnn_reshape(np.ones_like(samples))]) - normalize(ap_labels[sample_indices])) ** 2).mean(axis=None)

        # save mse on unlabeled data to analyze throughout training
        myfile = open(model_name+'_mse.txt', 'a')
        myfile.write("%s" % mse)
        myfile.write("\n")
        myfile.close()

        # Compute Time Elapsed
        elapsed_time = time.time() - start_time
        eta = elapsed_time
        print u'| TE: %s ' % time_format(eta),; sys.stdout.write(u'')

        #sys.stdout.flush()
        print('\n')

        # save models
        encoder_A.save('models/encoder_A_'+model_name+'_epoch_'+str(e)+'.h5')
        #predictor_A.save('models/predictor_A_'+model_name+'_epoch_'+str(e)+'.h5')
        decoder_A.save('models/decoder_A_'+model_name+'_epoch_'+str(e)+'.h5')
        encoder_B.save('models/encoder_B_'+model_name+'_epoch_'+str(e)+'.h5')
        decoder_B.save('models/decoder_B_'+model_name+'_epoch_'+str(e)+'.h5')

        test_sample_indices = np.random.choice(range(test_indices_range_ap[0],test_indices_range_ap[1]), 5, replace=False)
        sample_orig,_,_,_,_ = load_train_data_weighted_with_asset(data_file_l,data_file_u,indices_l=None,indices_u=test_sample_indices,
                                       unlabeled=True,labeled=False)
        zero_mask_test = create_zero_mask(sample_orig,0,3,1030)
        sample_orig = normalize_spec(sample_orig)
        tex_x_normed,test_x = gen_x_to_x_u.predict([cnn_reshape(sample_orig),cnn_reshape(zero_mask_test)])
        sample_orig = tex_x_normed.reshape(tex_x_normed.shape[0],tex_x_normed.shape[1])
        sample_orig_aug = sample_orig*zero_mask_test

        sample_diff = sample_orig-test_x.reshape(test_x.shape[0],test_x.shape[1])

        #save test results
        fig, axes = plt.subplots(20,1,figsize=(70, 20))
        for i in range(len(test_sample_indices)):
            axes[i*4].plot(sample_orig[i],c='r')
            axes[i*4].set_ylim((-4.,4))
            axes[1+4*i].plot(sample_orig_aug[i],c='g')
            axes[1+4*i].set_ylim((-4.,4))
            axes[2+4*i].plot(test_x[i],c='b')
            axes[2+4*i].set_ylim((-4.,4))
            axes[3+4*i].plot(sample_diff[i],c='m')
            axes[3+4*i].set_ylim((-0.3,0.3))

        plt.savefig('results/test_sample_ap_'+model_name+'_epoch_'+str(e)+'.jpg')

        test_sample_indices = np.random.choice(range(test_indices_range_synth[0],test_indices_range_synth[1]), 5, replace=False)
        _,_,_,sample_orig,sample_y = load_train_data_weighted_with_asset(data_file_l,data_file_u,indices_l=test_sample_indices,indices_u=None,
                                       unlabeled=False,labeled=True)
        zero_mask_test = create_zero_mask(sample_orig,0,3,1030)
        sample_orig = normalize_spec(sample_orig)
        tex_x_normed,test_x = gen_x_to_x_u.predict([cnn_reshape(sample_orig),cnn_reshape(zero_mask_test)])
        sample_orig = tex_x_normed.reshape(tex_x_normed.shape[0],tex_x_normed.shape[1])
        sample_orig_aug = sample_orig*zero_mask_test

        sample_diff = sample_orig-test_x.reshape(test_x.shape[0],test_x.shape[1])

        #save test results
        fig, axes = plt.subplots(20,1,figsize=(70, 20))
        for i in range(len(test_sample_indices)):
            axes[i*4].plot(sample_orig[i],c='r')
            axes[i*4].set_ylim((-4.,4))
            axes[1+4*i].plot(sample_orig_aug[i],c='g')
            axes[1+4*i].set_ylim((-4.,4))
            axes[2+4*i].plot(test_x[i],c='b')
            axes[2+4*i].set_ylim((-4.,4))
            axes[3+4*i].plot(sample_diff[i],c='m')
            axes[3+4*i].set_ylim((-0.3,0.3))

        plt.savefig('results/test_sample_synth_'+model_name+'_epoch_'+str(e)+'.jpg')

        plt.close('all')
```

```python
model_name='vae_ssl_cycle_test'
start_e = 0
```

# Train Models

```python
reporter=['l_loss','p_loss','u_loss']
epochs=200
batchsize=64
if start_e>0:
    start_e=start_e+1

data_file_l = 'ASSET_lower_temp.h5'
data_file_u = 'train_data.h5'

fit_model(model_name,data_file_l, data_file_u, epochs,reporter)
```

```python

```


# The Winding Number and the SSH model
### Christina Lee
### Category: Graduate
### Topological Physics Series
* [Quantum Anomolous Hall Effect and the Chern Number](../Graduate/Chern-Number.ipynb)
* [SSH Model and the Winding Number](../Graduate/Winding-Number.ipynb)

The Chern number isn't the only topological invariant.  We have multiple invariants, each convenient in their own situations.  The Chern number just happened to appear in one of the biggest, early examples, the Integer Quantum Hall Effect, but the winding number actually occurs much more often in a wider variety of circumstances.

How many times does the phase wrap as we transverse a closed loop?
$$
n = \frac{1}{2 \pi i } \oint \text{d}\phi =
\frac{1}{2\pi i } \oint \frac{\text{d}z}{z}
$$

This expression shows up in complex analysis with Residues and [the Cauchy Integral Formula](https://en.wikipedia.org/wiki/Cauchy%27s_integral_formula), but we're interested in applying this formula to topology.

## Topology and Homotopy

Topology is a general umbrella term for studying properties independent of deformation or coordinate systems.  If we go back to "what does topology formally mean?", it's a structure we can put on sets.  From there, we have a variety of different ways to study that structure, and one of those is Homotopy.

Homotopy considers two functions and asks whether or not they can be deformed into each other.

Here's a simple example:

We have positions $\vec{x}(t)$ over time with a fixed starting and stopping point.  And we have some fixed puncture point in space our function can never occupy.  We can classify all the possible paths by the number of times they go around the puncture point.

The domain in our situation is the unit circle $k$, and we want to know what the range looks like in terms of unit circles:
* Zero unit circles= a point?
* One Unit Circle?
* A Unit Circle followed by another Unit Circle?
* A Unit Circle, but flipped and traveled in the opposite direction?

Each of these is a different homotopy class.

## Su-Schrieffer-Heeger Model for Trans-polyacetylene

The Su-Schrieffer-Heeger Model for Trans-Polyacetylene hosts topological phases characterized by the winding number.

The chemical under doping has high electrical conductivity, opened the entire field of conductive polymers, and led to the 2000 Chemistry Nobel Prize [3].  To get to the model, we first need to look at the chemical structure,

Scary Organic Chemistry stuff... and due to my lack of understand Organic Chemistry stuff, I actually understood the physical situation wrong for a while, but now let me take a crack at explaining what I think is actually the case.

Though just the plain positions of the atoms are translationally symmetric with a period of one, due to [Pereils instability in 1D](https://en.wikipedia.org/wiki/Peierls_transition) the bonds break into having a translational period of 2.  This gives a degeneracy in the ground state.  The way the structure is tiled with "single-bond"--"double-bond" could just be shifted by one site, giving an equally viable chemical structure.

Now we have regions of uniform tiling and boundaries between them.  The boundaries then act like solitons with certain transitions probabilities to move around.  We are looking at the Hamiltonian that describes how this boundary soliton moves.

All we really need to know is we have a particle free to move in one dimension, with different transition probabilities depending on the type of bond.  For the technical history, see [4], [5], [6]. [6] makes the most sense.

$v$ and $w$ are our two transition probabilities, and we also have two different types of sites, $a$ and $b$. We can write down a hopping Hamiltonian from this information,
\begin{equation}
\mathcal{H}= v \sum_i \left(a_i^{\dagger} b_i + \text{h.c.} \right)
-w \sum_i \left(b_i^{\dagger} a_{i+1} + \text{h.c.} \right)
\end{equation}

We can Fourier transform the Hamiltonian,
\begin{equation}
\mathcal{\tilde{H}}= \sum_k v \left( a_k^{\dagger} b_k + b_k a_k^{\dagger} \right)
- w \left( e^{-i k}  b_k^{\dagger} a_k + e^{i k} a_k^{\dagger} b_k \right)
\end{equation}
and change forms
\begin{equation}
= \sum_k  \left(v -w \cos k \right) \left( b_k^{\dagger} a_k + a_k^{\dagger} b_k \right)
-i w \sin k \left( - b_k^{\dagger} a_k + a_k^{\dagger} b_k \right)
\end{equation}
\begin{equation}
=\begin{bmatrix}
a_k^{\dagger} & b_k^{\dagger}
\end{bmatrix}
\begin{bmatrix}
0 & \left(v-w \cos k\right) -i \left( -w \sin k \right) \\
\left(v-w \cos k\right) + i \left( -w \sin k \right) & 0 \\
\end{bmatrix}
\begin{bmatrix}
a_k \\ b_k
\end{bmatrix}
\end{equation}
to something where we can read off a quite useful form for this type of topological stuff:
\begin{equation}
\mathcal{\tilde{H}} = \vec{R} \cdot \vec{\sigma}
\end{equation}

\begin{equation}
R_x(k) = v - w \cos k \qquad \qquad R_y (k) = -w \sin k \qquad \qquad R_z = 0
\end{equation}

## Code up the Model

```julia
# Adding the Packages
using Plots
using LinearAlgebra
gr()
```

    Plots.GRBackend()

```julia
# Pauli Matrices
σx=[[0 1]
    [1 0]]
σy=[[0 -im]
    [im 0]]
σz=[[1 0]
    [0 -1]]
```

    2×2 Array{Int64,2}:
     1   0
     0  -1

```julia
# Functions
Rx(k::Float64,v=1,w=2)=v-w*cos(k)
Ry(k::Float64,v=1,w=2)=-w*sin(k)

R(k::Float64,v=1,w=2)=sqrt(Rx(k,v,w)^2+Ry(k,v,w)^2)

H(k::Float64,v=1,w=2)=Rx(k,v,w)*σx+Ry(k,v,w)*σy
```

    H (generic function with 3 methods)

```julia
# domain we will calculate on
l=314

ks=range(-π,stop=π,length=l)
dk=ks[2]-ks[1]
```

    0.02007407446383258

### Chiral Symmetry

A Hamiltonian is said to possess chiral symmetry if there exists a $U$ such that
$$
U H U^{-1} = -H  \qquad \qquad U U^{\dagger} =\mathbb{1}.
$$
Finding if $U$ even exists and determining its form if it exists is a problem for another time.  Today, multiple places said that $\sigma_z$ works for the SSH model, and we can confirm that it does.

A little less intellectually satisfying (at least for me), but it works.

We could test that equation analytically on pen and paper, analytically using 'SymPy', or by plugging in random 'k' values a bunch of times and assuming that's good enough.

I'm going the bunch of random k values route. Just keep evaluating the next cell till you're convinced.

Bonus Note:  We only have a situation with a winding number because we have chiral symmetry and have an odd number of dimensions.  If we have no chiral (and no other any other anti-unitary) symmetry, then we could only have the topologically trivial phase. That's why I'm making sure to mention this.  Check out the Periodic Table of Topological Insulators for more information.

```julia
k_test=rand()
σz*H(k_test)*σz^(-1)+H(k_test) # Should equal zero
```

    2×2 Array{Complex{Float64},2}:
     0.0+0.0im  0.0+0.0im
     0.0+0.0im  0.0+0.0im

### Homotopically Different Hamiltonians

We want to know if two sets of parameters $c_1=(v_1,w_1)$ and $c_2=(v_2,w_2)$ will describe topologically equivalent systems.

We are gifted by the fact have some convenient theorems relating the SPT (Symmetry Protected Topological) topology of the phases and band gap closings.

As we change the parameters, we will remain in the same topological phase as long as we don't close the gap or break symmetries.  [2]

The states are in the same phase if they can be smoothly deformed into each other.  In this context, smoothly deformed means connected by local unitary evolution. Nothing drastic is happening to the system.  Closing the gap is considered drastic.

Now we don't have to change the topological phase at a gap closing, but it's only possible there.

Beforehand doing a lot of calculations and work, if we can identify where band closings occur and regions where parameters can be perturbed and changed without causing band closings, we can reduce the number of things we need to solve later on.

Analytically, we know that the two eigenvalues occur at: (see QAHE post)
$$
R=\pm \sqrt{R_x^2+ R_y^2}
$$
$$
=\pm \sqrt{v^2+w^2 \cos^2 k -2 vw \cos k + w^2 \sin^2 k}
= \pm \sqrt{v^2 - 2 vw \cos k + w^2}
$$
The difference between the upper and lower band will be at its minimum when $\cos k$ is greatest,$k=0$.
$$
=\pm \sqrt{(v-w)^2}
$$

So when $v=w$, the gap closes.  This $v=w$ line in parameter space could separate two different topological phases.  Now we need to perform some calculations to see if that is they are actually different phases.

To more quickly see which side of the dividing line a parameter set falls on, I'm instead going to write out parameters in terms of $v$ and $d = v-w$.  This way, the sign of $d$ can quickly tell me which phase we are in.

If d is positive, we are in the Purple phase, designated so because that's what I am using for my color scheme. The Purple phase also turns out to be the topological phase, as we will see later.

When d is negative, we are in the Turquoise phase, again because of my color scheme.  This phase is topologically trivial.

```julia
# Parameters chosen to look at
va=[1.0, 0.5,1.0,  0.0,0.5,  0.4,0.6]
da=[0.0, 0.5,-0.5, 0.5,-0.5, 0.2,-0.2]

# w values corresponding to chosen ds
wa=round.(va+da;sigdigits=2) #Floating point error was making it print bad

# how to plot chosen parameters
colors=[colorant"#aa8e39",
    colorant"#592a71",colorant"#4e918f",
    colorant"#310c43",colorant"#226764",
    colorant"#cca7df",colorant"#a0d9d7",
    ]
styles=[:solid,
    :dash,:dash,
    :solid,:solid,
    :dot,:dot]
widths=[10,15,5,15,5,10,3];
```

## Band diagrams for different parameters

When $d=0$, we have the gold line with zero band-gap.

If $d\neq 0$, then the band gap is not zero.

If $v$ and $w$ are flipped, the model will look identical in its energy structure.  We can see the two different sets of purple and turquoise lines plotted over each other.  Only when we look at the phase of the wavefunctions can we see that flipping the $v$ and $w$ actually does have quite an influence on the solution.

```julia
plot()
for ii in 1:length(va)
    plot!(ks,R.(ks,va[ii],wa[ii])
        ,label="v=$(va[ii]) w=$(wa[ii])"
        ,linewidth=widths[ii],color=colors[ii],linestyle=styles[ii])

    plot!(ks,-R.(ks,va[ii],wa[ii])
        ,label=""
        ,linewidth=widths[ii],color=colors[ii],linestyle=styles[ii])
end
plot!(title="Band diagrams for different parameters",
xlabel="Momentum",ylabel="Energy")
```

## Homotopy of Hamiltonian Vector

We can look at either the homotopy of the Hamiltonian or the homotopy of the eigenfunctions.

Looking at the Hamiltonian seems easier since we don't have to go through the work of calculating the wavefunctions, especially if we have a complicated system, but homotopy is a geometric, almost pictorial thing. How do we go about getting something like that for an operator?

Let's go back to how we wrote our Hamiltonian down, both this one and the QAHE one before,
$$
\mathcal{H}=\vec{R}(k) \cdot \vec{\sigma}.
$$
Here we have a 1-1 correspondence between the Hamiltonian and a geometric object, this $\vec{R}$ vector.  When we look at how it depends on $k$, we get insight into how $\mathcal{H}$ depends on $k$ as well.

The two different groups, purple and turquoise, will have two different behaviors.  $\vec{R}(k)$ for purple will circle the origin like $S^1$ the unit circle, whereas $\vec{R}(k)$ for turquoise will not circle the origin and will not be like $S^1$.

```julia
plot()
for ii in 1:length(va)
    plot!(Rx.(ks,va[ii],wa[ii]),
        Ry.(ks,va[ii],wa[ii])
    ,label="v=$(va[ii]) , d=$(da[ii])"
    ,linewidth=5,color=colors[ii],linestyle=styles[ii])
end

statval=5
scatter!(Rx.(ks,va[statval],wa[statval]),Ry.(ks,va[statval],wa[statval])
    ,label="",markersize=10,color=colors[statval])

scatter!([0],[0],label="Origin",
        markersize=20,markershape=:star5,color=colorant"#6f4f0d")

plot!(title="R(k) for different parmeters",
xlabel="R_x", ylabel="R_y",legend=:bottomright,aspect_ratio=1)
```

## Wavefunction

Going back to what we did in the QAHE post again, we have an analytical expression for the wavefunction:
\begin{equation}
\Psi = \frac{1}{\sqrt{2}}
\begin{bmatrix}
-1\\
\frac{R_x + i R_y}{R}
\end{bmatrix}
\end{equation}

Now that we don't have $R_z$, the magnitudes of both components are constant and uniform. For the first component, the phase also remains constant, but the phase of the second component varies.  It's the behavior of this phase that we will be looking at, and that will determine whether or not the system is topological.

```julia
um1=-1/sqrt(2)

function um2(k::Float64,v=1,w=2)
    return 1/(sqrt(2)*R(k,v,w))*(Rx(k,v,w)+im*Ry(k,v,w))
end
```

    um2 (generic function with 3 methods)

## Plotting the Phase

In plotting the phase for the different parameter combinations, we can really see the differences between the topological phases.  In the turquoise group that didn't encircle zero, the phase changes sinusoidally, going up then back down again, so on and so forth around zero.

But for our purple states $d>0$, the phase just keeps increasing, so we get jumps as we confine it between $-\pi$ and $\pi$. The phase itself is continuous; it just goes across a branch cut which gives us a discontinuity in how we write it down.

As for the $d=0$ systems, those are both boundary cases with more complicated behavior.

```julia
plot()
for ii in 1:length(va)
    plot!(ks,angle.(um2.(ks,va[ii],wa[ii]))
    ,label="v=$(va[ii]) , d=$(da[ii])",linewidth=5
    ,color=colors[ii],linestyle=styles[ii])
end
plot!(title="Phase",xlabel="k",ylabel="angle")
```

We can look at the effect of our decision of how to take an angle by rotating the system before applying the $-\pi$-$\pi$ boundary.

If we rotate the system by $\pi/4$ first, the discontinuity in the topological wavefunctions occurs at a different k-location, but it doesn't go away.  This same thing happened in the QAHE Chern number situation.  We can write something different and make a problem area occur in a different spot, but it's still going to occur somewhere.  We can't get rid of the wrapping behavior of the topological systems by any amount of looking at it differently or smooth manipulations.  We can only move the discontinuity that arises from it to a different location.

```julia
plot()
for ii in 1:length(va)
    plot!(ks,angle.(exp(im*π/4)*um2.(ks,va[ii],wa[ii]))
    ,label="v=$(va[ii]) , w=$(wa[ii])",linewidth=5
    ,color=colors[ii])
end
plot!(title="Phase",xlabel="k",ylabel="angle")
```

We've qualitatively seen the difference between the phases, but now let's quantitatively look at the difference between the phases. This is as simple integrating the formula in the introduction,
\begin{equation}
n= \oint \frac{\text{d} z}{z}
\end{equation}

```julia
function Winding_phi(k,v,w)
    dum2=(um2.(k[2:end],v,w).-um2.(k[1:(end-1)],v,w))
    return 1/(2π*im)*sum(dum2./um2.(k[2:end],v,w) )
end
```

    Winding_phi (generic function with 1 method)

Calculating the winding number specifically for the parameter combinations we've been looking at so far, we can see that the positive phase $d>0$ has a winding number of 1 and the negative phase $d<0$ has a winding number of 0.

```julia
println("|Phase \t n \t| d \t v \t w |\t Real \t Imag")
for ii in 1:length(va)
    temp=Winding_phi(ks,va[ii],wa[ii])

    println("| ",sign(da[ii]),"\t",round(real(temp),digits=1),"\t|",
        da[ii],"\t",va[ii],"\t",wa[ii],"|\t",
        round(real(temp),digits=5),"\t",round(imag(temp),digits=5))
end
```

    |Phase 	 n 	| d 	 v 	 w |	 Real 	 Imag
    | 0.0	0.5	|0.0	1.0	1.0|	0.4968	-0.3208
    | 1.0	1.0	|0.5	0.5	1.0|	0.99989	-0.01171
    | -1.0	0.0	|-0.5	1.0	0.5|	1.0e-5	-0.00167
    | 1.0	1.0	|0.5	0.0	0.5|	0.99993	-0.01004
    | -1.0	0.0	|-0.5	0.5	0.0|	0.0	0.0
    | 1.0	1.0	|0.2	0.4	0.6|	0.99982	-0.01405
    | -1.0	0.0	|-0.2	0.6	0.4|	3.0e-5	-0.00401

But we can also calculate the winding number for the entire grid of parameter values. Here we can much more obviously see how $d=0 \rightarrow v=w$ represents a phase transition between two different topological phases.

```julia
vaa=repeat(range(0,stop=1,length=100),1,100)
waa=transpose(vaa)

ϕaa=zeros(Complex{Float64},100,100)
for ii in 1:100
    for jj in 1:100
        ϕaa[ii,jj]=Winding_phi(ks,vaa[ii,jj],waa[ii,jj])
    end
end
```

```julia
heatmap(vaa[:,1],waa[1,:],real.(ϕaa))
plot!(xlabel="v",ylabel="w", title="Two Different Topological Phases")
```

## Conclusion

Systems in one dimension with chiral symmetry can host topological phases characterized by the winding number.

The winding behavior appears in both the Hamiltonian and the wavefunctions.

Transitions between topological phases occur at band gap closings.

The Su-Schrieffer-Heeger model exhibits two different topological phases.

[1] Public Domain, https://commons.wikimedia.org/w/index.php?curid=1499462

[2] Chen, Xie, Zheng-Cheng Gu, and Xiao-Gang Wen. "Local unitary transformation, long-range quantum entanglement, wave function renormalization, and topological order." Physical review b 82.15 (2010): 155138. https://arxiv.org/pdf/1004.3835.pdf

[3] https://www.nobelprize.org/prizes/chemistry/2000/popular-information/

[4] Su, W-P_, J. R. Schrieffer, and A. J. Heeger. "Soliton excitations in polyacetylene." Physical Review B 22.4 (1980): 2099.

[5] Heeger, Alan J., et al. "Solitons in conducting polymers." Reviews of Modern Physics 60.3 (1988): 781.

[6] Takayama, Hajime, Yo R. Lin-Liu, and Kazumi Maki. "Continuum model for solitons in polyacetylene." Physical Review B 21.6 (1980): 2388.

[7] Meier, E. J. et al. Observation of the topological soliton state in the Su-Schrieffer-Heeger model. Nat. Commun. 7, 13986 doi: 10.1038/ncomms13986 (2016)

http://paletton.com/#uid=31b0J0kllll8rOUeTt+rNcHBJ42

```julia

```


# Re-creating [Capillary Hysteresis in Neutrally Wettable Fibrous Media: A Pore Network Study of a Fuel Cell Electrode](http://link.springer.com/10.1007/s11242-017-0973-2)
# Part C: Purcell Meniscus Model

## Introduction

In the final part of this series we take a deeper look at the Purcell meniscus model, which is central to the whole study. The fundamental concept is that the fibers form throats with properties similar to that of a torus:

As the fluid invades throaugh the center of the torus the meniscus is pinned to the surface and the "effective" contact angle becomes influenced by the converging diverging geometry and is a function of the filling angle $\alpha$. The shape of the meniscus as the invading phase moves upwards through the torus with key model parameters is shown below.

Different intrinsic contact angles through invading phase: (a) 60$^\circ$, (b) 90$^\circ$ and (c) 120$^\circ$. All scenarios clearly show an inflection of the meniscus curvature signifying a switch in the sign of the capillary pressure from negative to positive. This inflection is predicted to occur for all contact angles by the model with varying filling angle. The capillary pressure can be shown to be:

$P_C = -2\sigma cos(\theta-\alpha))/(r+R(1-cos(\alpha))$

## Set up
We will set up a trivially small network with one throat to demonstrate the use of the meniscus model. Here we do the imports and define a few functions for plotting.

```python
import matplotlib.pyplot as plt
import numpy as np
import sympy as syp
from ipywidgets import interact
%pylab inline
pylab.rcParams['figure.figsize'] = (5, 5)

theta = 90
fiberRad = 5e-6
throatRad = 2e-5
max_bulge = 1e-5

# Parameters for plotting fibers
x, R, rt, s, t = syp.symbols('x, R, rt, s, t')
points = np.linspace(-fiberRad, fiberRad, 1001)
y = R*syp.sqrt(1- (x/R)**2)
r = rt + (R-y)
rx = syp.lambdify((x, R, rt), r, 'numpy')

def plot_arc(ax, arc_cen, arc_rad, arc_angle, num_points=100):
    r'''
    plot the meniscus arc from neg to pos angle
    '''
    angles = np.linspace(-arc_angle, arc_angle, num_points)
    x = arc_cen - arc_rad*np.cos(angles)
    y = arc_rad*np.sin(angles)
    print(x)
    print(y)
    ax.plot(x, y, 'b-')
```

    Populating the interactive namespace from numpy and matplotlib

Now we define our two pore network and add the meniscus model in several modes: 'max' returns the maximum pressure experienced by the meniscus as it transitions through the throat, i.e. the burst entry pressure. 'touch' is the pressure at which the meniscus has protruded past the throat center a distance defined by the 'touch_length' dictionary key. In network simulations this could be set to the pore_diameter. Finally the 'men' mode accepts a target_Pc parameter and returns all the mensicus information required for assessing cooperative filling or plotting.

```python
import openpnm as op
import openpnm.models.physics as pm
net = op.network.Cubic(shape=[2, 1, 1], spacing=5e-5)
geo = op.geometry.StickAndBall(network=net,
                               pores=net.pores(),
                               throats=net.throats())
phase = op.phases.Water(network=net)
phase['pore.contact_angle'] = theta
phys = op.physics.Standard(network=net,
                           phase=phase,
                           geometry=geo)
geo['throat.diameter'] = throatRad*2
geo['throat.touch_length'] = max_bulge
phys.add_model(propname='throat.max',
               model=pm.meniscus.purcell,
               mode='max',
               r_toroid=fiberRad)
phys.add_model(propname='throat.touch',
               model=pm.meniscus.purcell,
               mode='touch',
               r_toroid=fiberRad)
phys.add_model(propname='throat.meniscus',
               model=pm.meniscus.purcell,
               mode='men',
               r_toroid=fiberRad,
               target_Pc=1000)
touch_Pc = phys['throat.touch'][0]
print(touch_Pc)
max_Pc = phys['throat.max'][0]
print(max_Pc)
```

    4114.018827899374
    5878.390833049412

We define a plotting function that uses the meniscus data:
$\alpha$ is filling angle as defined above,
$radius$ is the radius of curvature of the mensicus,
$center$ is the position of the centre of curvature relative to the throat center along the axis of the throat,
$\gamma$ is the angle between the throat axis and the line joining the meniscus center and meniscus contact point.

```python
def plot_phys(target_Pc, print_data=False):
    phys.add_model(propname='throat.meniscus',
               model=pm.meniscus.purcell,
               mode='men',
               r_toroid=fiberRad,
               target_Pc=target_Pc)
    throatRad = geo['throat.diameter'][0]/2
    theta = np.deg2rad(phys['pore.contact_angle'][0])
    men_data = {}
    men_data['alpha'] = phys['throat.meniscus.alpha']
    men_data['gamma'] = phys['throat.meniscus.gamma']
    men_data['radius'] = phys['throat.meniscus.radius']
    men_data['center'] = phys['throat.meniscus.center']
    fig, ax = plt.subplots()
    ax.plot(points, rx(points, fiberRad, throatRad), 'k-')
    ax.plot(points, -rx(points, fiberRad, throatRad), 'k-')
    arc_cen = men_data['center']
    arc_rad = men_data['radius']
    arc_angle = men_data['gamma']
    angles = np.linspace(-arc_angle, arc_angle, 100)
    arcx = arc_cen - arc_rad*np.cos(angles)
    arcy = arc_rad*np.sin(angles)
    ax.plot(arcx, arcy, 'b-')
    ax.axis('equal')
    ax.ticklabel_format(style='sci', axis='both', scilimits=(-6,-6))
    if print_data:
        print(men_data)
    return ax
```

We can see that the touch_Pc calculated earlier, corresponds with the tip of the meniscus exceeding the max_bulge parameter. Try changing this and re-running to see what happens.

```python
ax = plot_phys(target_Pc=touch_Pc, print_data=True)
ax.plot([max_bulge, max_bulge], [-throatRad, throatRad], 'r--')
```

We can interact with the mensicus model by changing the target_Pc parameter.

```python
interact(plot_phys, target_Pc=(-max_Pc, max_Pc, 10))
```

    A Jupyter Widget

    <function __main__.plot_phys>


# Unobserved Heterogeneity and Finite Mixture Models

Unobserved heterogeneity is a concern in every econometric application. Keane and Wolpin (1997) face the problem that individuals at the age of sixteen report varying years of schooling. Neglecting the issue of measurement error, it is unlikely that the differences in initial schooling are caused by exogenous factors. Instead, the schooling decision is affected by a variety of endogenous factors such as parental investement, school and teacher quality, intrinsic motivation, and ability. Without correction, estimation methods fail to recover the true parameters.

One solution would be to extend the model and incorporate the whole human capital investement process up to the age where initial schooling was zero. Although such a model would be extremely interesting, it is also almost infeasible to model that many factors in terms of modeling, computation and data.

Another solution is to employ individual fixed-effects. Then, the state space comprises a dimension with has the same number of unique values as there are individuals in the sample. Thus, you have to compute the decision rules for every individual for the whole state space separately which is computationally infeasible.

Keane and Wolpin (1997) resort to model unobserved heterogeneity with a finite mixture. A mixture model can be used to model the presence of subpopulations (types) in the general population without requiring the observed data to identify the affiliation to a group. In contrast to fixed-effects, the number of subpopulations is much lower than the number of individuals. There is also no fixed and unique assignment to one subpopulation, but relations are defined by a probability mass function.

Each type has a preference for a particular choice which is modeled by a constant in the utility functions. For working alternatives, $w$, the constant is in the log wage equation whereas for non-working alternatives, $n$, it is in the nonpecuniary reward. Note that **respy** allows for type-specific effects in every utility component. Keane and Wolpin (1997) call it endowment with the symbol $e_{ak}$ for type $k$ and alternative $a$.

$$\begin{align}
    \log(W(s_t, a_t)) = x^w\beta^w  + e_{ak} + \epsilon_{at}\\
    N^n(s_t, a_t) = x^n\beta^n + e_{ak} + \epsilon_{at}
\end{align}$$

To estimate model parameters with maximum likelihood, the likelihood contribution for one individual is defined as the joint probability of choices and wages accumulated over time.

$$
    P(\{a_t\}^T_{t=0} \mid s^-_t, e_{ak}, W_t) =
    \prod^T_{t = 0} p(a_t, \mid s^-_t, e_{ak}, W_t)
$$

We can weight the contribution for type $k$ with the probability for being the same type to get the unconditioned likelihood contribution of an individual.

$$
    P(\{a_t, W_t\}^T_{t=0}) = \sum^K_{k=1} \pi_k
        P(\{a_t\}^T_{t=0} \mid s^-_t, e_{ak}, W_t)
$$

To avoid misspecification of the likelihood, $\pi_k$ must be a function of all individual characteristics which are determined before individuals enter the model horizon and are not the result of exogenous factors. The type-specific probability $\pi_k = f(x^\pi \beta^\pi)$ is calculated with softmax function based on a vector of covariates $x^\pi$ and a matrix of coefficients $\beta^\pi$ for each type-covariate combination.

$$
    \pi_k = f(x^\pi \beta^\pi_k) =
        \frac{\exp{\{x^\pi \beta^\pi_k\}}}{\sum^K_{k=1} \exp \{x^\pi \beta^\pi_k\}}
$$

To implement a finite mixture, we have to express $e_{ak}$ and $\beta^\pi$ in the parameters. As an example, we start with the basic Robinson Crusoe Economy.

```python
import io
import pandas as pd
import respy as rp
```

```python
params, options = rp.get_example_model("robinson_crusoe_basic", with_data=False)
params
```

      value

      category
      name

      delta
      delta
      0.95

      wage_fishing
      exp_fishing
      0.30

      nonpec_fishing
      constant
      -0.20

      nonpec_hammock
      constant
      2.00

      shocks_sdcorr
      sd_fishing
      0.50

      sd_hammock
      0.50

      corr_hammock_fishing
      0.00

We extend the model by allowing for different periods of experience in fishing at $t = 0$. Robinsons starts with zero, one or two experience in fishing because of different tastes for fishing.

```python
initial_exp_fishing = """
category,name,value
initial_exp_fishing_0,probability,0.33
initial_exp_fishing_1,probability,0.33
initial_exp_fishing_2,probability,0.34
"""
```

```python
initial_exp_fishing = pd.read_csv(io.StringIO(initial_exp_fishing), index_col=["category", "name"])
initial_exp_fishing
```

      value

      category
      name

      initial_exp_fishing_0
      probability
      0.33

      initial_exp_fishing_1
      probability
      0.33

      initial_exp_fishing_2
      probability
      0.34

In the next step, we add type-specific endowment effects $e_{ak}$. We assume that there exist three types and the additional utility is increasing from the first to the third type. For computational simplicity, the benefit of the first type is normalized to zero such that all other types are in relation to the first.

```python
endowments = """
category,name,value
wage_fishing,type_1,0.2
wage_fishing,type_2,0.4
"""
```

```python
endowments = pd.read_csv(io.StringIO(endowments), index_col=["category", "name"])
endowments
```

      value

      category
      name

      wage_fishing
      type_1
      0.2

      type_2
      0.4

We assume no effect for choosing the hammock.

At last, we need to specify the probability mass function which relates individuals to types. We simply assume that initial experience is positively correlated with a stronger taste for fishing. For a comprehensive overview on how to specify distributions with multinomial coefficients, see the guide on the [initial conditions](initial_conditions.ipynb). Note that, the distribution is also only specified for type 1 and 2 and the coefficients for type 1 are left out for a parsimonuous representation. You cannot use probabilities as type assignment cannot be completely random. The following example is designed to specify a certain distribution and recover the pattern in the data. In reality, the distribution of unobservables is unknown.

First, we define that Robinsons without prior experience are of type 0. Thus, we make the coefficients for type 1 and 2 extremely small. Robinsons with one prior experience are of type 1 with probability 0.66 and type 2 with 0.33. For two periods of experience for fishing, the share of type 1 individuals is 0.33 and of type 2 is 0.66. The coefficients for type 1 and 2 are simply the log of the probabilities.

At last, we add a sufficiently large integer to all coefficients. The coefficient of type 0 is implicitly set to zero, so the distribution samples type 0 individuals for one or two experience in fishing. By shifting the parameters with a positive value, this is prevented. At the same time, the softmax function is shift-invariant and the relation of type 1 and type 2 shares is preserved.

```python
type_probabilities = """
category,name,value
type_1,initial_exp_fishing_0,-100
type_1,initial_exp_fishing_1,-0.4055
type_1,initial_exp_fishing_2,-1.0986
type_2,initial_exp_fishing_0,-100
type_2,initial_exp_fishing_1,-1.0986
type_2,initial_exp_fishing_2,-0.4055
"""
```

```python
type_probabilities = pd.read_csv(io.StringIO(type_probabilities), index_col=["category", "name"])
type_probabilities += 10
type_probabilities
```

      value

      category
      name

      type_1
      initial_exp_fishing_0
      -90.0000

      initial_exp_fishing_1
      9.5945

      initial_exp_fishing_2
      8.9014

      type_2
      initial_exp_fishing_0
      -90.0000

      initial_exp_fishing_1
      8.9014

      initial_exp_fishing_2
      9.5945

The covariates used for the probabilities are defined below.

```python
type_covariates = {
    "initial_exp_fishing_0": "exp_fishing == 0",
    "initial_exp_fishing_1": "exp_fishing == 1",
    "initial_exp_fishing_2": "exp_fishing == 2",
}
type_covariates
```

    {'initial_exp_fishing_0': 'exp_fishing == 0',
     'initial_exp_fishing_1': 'exp_fishing == 1',
     'initial_exp_fishing_2': 'exp_fishing == 2'}

In the next step, we put all pieces together to get the complete model specification.

```python
params = params.append([initial_exp_fishing, endowments, type_probabilities])
params
```

      value

      category
      name

      delta
      delta
      0.9500

      wage_fishing
      exp_fishing
      0.3000

      nonpec_fishing
      constant
      -0.2000

      nonpec_hammock
      constant
      2.0000

      shocks_sdcorr
      sd_fishing
      0.5000

      sd_hammock
      0.5000

      corr_hammock_fishing
      0.0000

      initial_exp_fishing_0
      probability
      0.3300

      initial_exp_fishing_1
      probability
      0.3300

      initial_exp_fishing_2
      probability
      0.3400

      wage_fishing
      type_1
      0.2000

      type_2
      0.4000

      type_1
      initial_exp_fishing_0
      -90.0000

      initial_exp_fishing_1
      9.5945

      initial_exp_fishing_2
      8.9014

      type_2
      initial_exp_fishing_0
      -90.0000

      initial_exp_fishing_1
      8.9014

      initial_exp_fishing_2
      9.5945

```python
options["covariates"] = {**options["covariates"], **type_covariates}
options["simulation_agents"] = 10_000
options
```

    {'solution_draws': 100,
     'solution_seed': 456,
     'n_periods': 5,
     'simulation_agents': 10000,
     'simulation_seed': 132,
     'estimation_draws': 100,
     'estimation_seed': 100,
     'estimation_tau': 0.001,
     'interpolation_points': -1,
     'covariates': {'constant': '1',
      'initial_exp_fishing_0': 'exp_fishing == 0',
      'initial_exp_fishing_1': 'exp_fishing == 1',
      'initial_exp_fishing_2': 'exp_fishing == 2'}}

Let us simulate a dataset to see whether the distribution of types can be recovered from the data.

```python
simulate = rp.get_simulate_func(params, options)
df = simulate(params)
```

```python
df.query("Period == 0").groupby("Experience_Fishing").Type.value_counts(normalize="rows").unstack().fillna(0)
```

      Type
      0
      1
      2

      Experience_Fishing

      0
      1.000000
      0.000000
      0.000000

      1
      0.000000
      0.665548
      0.334452

      2
      0.000296
      0.330278
      0.669426

We also know that type 1 and 2 experience a higher utility for choosing fishing. Here are the choice probabilities for each type.

```python
df.groupby("Type").Choice.value_counts(normalize=True).unstack()
```

      Choice
      fishing
      hammock

      Type

      0
      0.426571
      0.573429

      1
      0.992602
      0.007398

      2
      0.998036
      0.001964


## Causal Analysis in Settings where the control group is orders of magnitude larger than the treatment group

```python
%load_ext autoreload
%autoreload 2

```

```python
import numpy as np
import pandas as pd

#dowhy
import dowhy
from dowhy import CausalModel
import dowhy.datasets, dowhy.plotter

#econml-scikit-learn
import econml
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LassoCV
from sklearn.ensemble import GradientBoostingRegressor

```

#### Filter Out Unnecessary Warnings

```python
import logging
logging.getLogger("dowhy").setLevel(logging.WARNING)
import warnings
from sklearn.exceptions import DataConversionWarning
warnings.filterwarnings(action='ignore', category=DataConversionWarning)

```

```python
np.random.seed(seed = 0)
```

**For the clustering** we use as means $\mu_1, \mu_2$. We choose $w$ to be parallel in the line connecting the two means.
 $\begin{align}w=\frac{\mu_2-\mu_1}{||\mu_2-\mu_1||}\end{align}$

The bias $b$ can be adjusted such that the hyperplane passes through a point connecting the two means:
$
    \begin{align}
    b = \frac{(\mu_2-\mu_1)}{||\mu_2-\mu_1||}^T(\mu_1 + \lambda (\mu_2-\mu_1))
    \end{align}$

```python
def generate_state(X, w = None, b = None, cluster = False):
    """
    Separates Patients
    with a linear hyperplane
    Patients above the hyperplane are healthy
    and below are sick

    """
    dims = X.shape[1]
    if not cluster:
        w = np.random.rand(dims, 1)
        w = w/np.sum(w)
        b = 1

    S = ((X@w-b) > 0).astype(int)

    return S

def generate_X(dims:int = 5, N:int = 10**4, cluster = False):

    """
    Generates X from a multidimension Normal distribution
    with the identity matrix for covariance

    """
    if not cluster:
        mu  = np.zeros(dims)
        loc = np.eye(dims)
        X = np.random.multivariate_normal(mean = mu, cov = loc, size = N)
        w = None
        b = None

    else:
        mu1 = np.zeros(dims)
        mu2 = np.ones(dims)
        loc = np.eye(dims)
        N1 = int(N*0.97)
        N2 = N-N1
        X1 = np.random.multivariate_normal(mean = mu1, cov = loc, size = N1)
        X2 = np.random.multivariate_normal(mean = mu2, cov = loc, size = N2)
        w = np.expand_dims((mu2-mu1)/np.sum(mu2-mu1), axis = 1)
        l = 0.85
        b = w.T@((1-l)*mu1[:,np.newaxis] + l*(mu2[:,np.newaxis]))
        X = np.concatenate((X1, X2), axis = 0)
        np.random.shuffle(X)
        print("{} people from cluster 1 generated".format(N1))
        print("{} people from cluster 2 generated".format(N2))

    return X, w, b

def treatment_assignment(S, M):
         """
         assign sick in treatment with p1
         assing healthy in treatment with p2
         S:State
         M:Number of people in treatment group
         """
         #generate a uniform sample
         uniform = np.random.uniform(0,1,len(S))

         helper  = np.zeros_like(uniform)

         #assign healthy people to pseudo treatment
         #with probability 0.95
         helper[(uniform <= 0.98) & (S[:,0] == 1)] = 1

         #assign sick people to pseudo treatment
         #with probability 0.2
         helper[(uniform >= 0.95) & (S[:,0] == 0)]  = 1

         T = np.zeros_like(helper)
         index = np.where(helper == 1)[0]
         choose = np.random.choice(index, size = M, replace = False)

         #take a random sample of M from the pseudo treatment
         # and assign it to real treatment
         T[choose] = 1
         T =np.expand_dims(T, 1).astype(int)
         T  = np.where(T==1, True, False)
         return T

def make_experiment(dims, N, M, cluster = False):

    X, w, b = generate_X(dims, N, cluster = cluster)
    S = generate_state(X, w = w, b = b, cluster = cluster)
    T = treatment_assignment(S,M)

    print_stats(S,T)

    Yf, Ycf, Y = create_outcome(S,T)

    data = create_pandas(X, S, T, Yf, Ycf, Y)

    return X, S, T, data

def create_pandas(X, S, T, Yf, Ycf, Y):

    columns = ['f'+str(i) for i in range(X.shape[1])]
    columns.extend(['S', 'Tr', 'Yf', 'Ycf', 'Y'])

    data = pd.DataFrame(np.concatenate([X,S,T, Yf, Ycf, Y], axis = 1), columns = columns)
    return data

def print_stats(S,T):
    print("Population Size:", len(S))
    print("Sick Population Size:", (S==1).sum())
    print("Treatment group Size:", T.sum())
    print("Sick People in Treatment group:", ((S==1)&(T==1)).sum())

def create_outcome(S,T):

    Yf  = np.ones_like(S)
    Ycf = np.ones_like(S)
    Ycf[(S==1)] = 0
    Y = np.ones_like(S)
    Y[(S==1)&(T==0)] = 0

    return Yf, Ycf, Y
```

```python
dims = 5
X, S, T, data= make_experiment(dims, 10**6, 1000, cluster=True)
data.Tr = data.Tr.astype(bool)
```

    970000 people from cluster 1 generated
    30000 people from cluster 2 generated
    Population Size: 1000000
    Sick Population Size: 46787
    Treatment group Size: 1000
    Sick People in Treatment group: 465

```python
data.head()
```

      f0
      f1
      f2
      f3
      f4
      S
      Tr
      Yf
      Ycf
      Y

      0
      0.900766
      -0.040178
      -1.083424
      0.121346
      -0.247905
      0.0
      False
      1.0
      1.0
      1.0

      1
      0.178655
      -0.498405
      1.631524
      1.267724
      -0.245090
      0.0
      False
      1.0
      1.0
      1.0

      2
      -1.532263
      -1.673195
      -0.825910
      -2.163854
      -1.181748
      0.0
      False
      1.0
      1.0
      1.0

      3
      0.598833
      -0.745410
      0.673457
      -0.178242
      0.286164
      0.0
      False
      1.0
      1.0
      1.0

      4
      1.044861
      -1.624495
      -1.035050
      0.136189
      0.205054
      0.0
      False
      1.0
      1.0
      1.0

```python
common_causes = ['f'+str(i) for i in range(dims)]
#common_causes.append('S')
data_dict = {'df':data,
             'treatment_name': 'Tr',
             'outcome_name': 'Y',
             'common_causes_names': common_causes,
             'time_val': None,
             'instrument_names': None,
             'dot_graph': None,
             'gml_graph': None,
             'ate': None}

```

```python
model = CausalModel(
        data=data_dict['df'],
        treatment=data_dict["treatment_name"],
        outcome=data_dict["outcome_name"],
        common_causes=data_dict["common_causes_names"],
        )
```

    WARNING:dowhy.causal_model:Causal Graph not provided. DoWhy will construct a graph based on data inputs.

```python
identified_estimand = model.identify_effect(proceed_when_unidentifiable=True)

```

    WARNING:dowhy.causal_identifier:If this is observed data (not from a randomized experiment), there might always be missing confounders. Causal effect cannot be identified perfectly.

```python
estimate = model.estimate_effect(identified_estimand,
                                 method_name="backdoor.propensity_score_stratification",
                                target_units='ate',
                                test_significance = None,
                                evaluate_effect_strength = False,
 confidence_intervals=False)
print(estimate)
print("Causal Estimate is " + str(estimate.value))
```

    *** Causal Estimate ***

    ## Identified estimand
    Estimand type: nonparametric-ate
    ### Estimand : 1
    Estimand name: backdoor
    Estimand expression:
      d
    ─────(Expectation(Y|f3,f4,f0,f1,f2))
    d[Tr]
    Estimand assumption 1, Unconfoundedness: If U→{Tr} and U→Y then P(Y|Tr,f3,f4,f0,f1,f2,U) = P(Y|Tr,f3,f4,f0,f1,f2)
    ### Estimand : 2
    Estimand name: iv
    No such variable found!

    ## Realized estimand
    b: Y~Tr+f3+f4+f0+f1+f2
    Target units: ate

    ## Estimate
    Mean value: 0.09541018707595271

    Causal Estimate is 0.09541018707595271

```python
data.Yf.mean() - data.Ycf.mean()
```

```python
data[data.Tr==1].Y.mean() - data[data.Tr == 0].Y.mean()
```

```python
data[(data.S==1) & (data.Tr ==1)].Y.mean() - data[(data.S==1) & (data.Tr ==0)].Y.mean()
```

```python
data[(data.S==0) & (data.Tr ==1)].Y.mean() - data[(data.S==0) & (data.Tr ==0)].Y.mean()
```

```python
model.view_model()
```

```python
data.head()
```

      f0
      f1
      f2
      f3
      f4
      S
      Tr
      Yf
      Ycf
      Y
      propensity_score
      strata
      dbar
      d_y
      dbar_y

      0
      0.900766
      -0.040178
      -1.083424
      0.121346
      -0.247905
      0.0
      False
      1.0
      1.0
      1.0
      0.000556
      21.0
      1
      0.0
      1.0

      1
      0.178655
      -0.498405
      1.631524
      1.267724
      -0.245090
      0.0
      False
      1.0
      1.0
      1.0
      0.001522
      42.0
      1
      0.0
      1.0

      2
      -1.532263
      -1.673195
      -0.825910
      -2.163854
      -1.181748
      0.0
      False
      1.0
      1.0
      1.0
      0.000049
      0.0
      1
      0.0
      1.0

      3
      0.598833
      -0.745410
      0.673457
      -0.178242
      0.286164
      0.0
      False
      1.0
      1.0
      1.0
      0.000833
      31.0
      1
      0.0
      1.0

      4
      1.044861
      -1.624495
      -1.035050
      0.136189
      0.205054
      0.0
      False
      1.0
      1.0
      1.0
      0.000429
      15.0
      1
      0.0
      1.0

## Using Econ ML

dml_estimate = model.estimate_effect(identified_estimand,
                                    method_name="backdoor.econml.dml.DMLCateEstimator",
                                    method_params={
                                        'init_params': {'model_y':GradientBoostingRegressor(),
                                                        'model_t': GradientBoostingRegressor(),
                                                        'model_final':LassoCV(fit_intercept=False), },
                                        'fit_params': {}
                                     })

dml_estimate.params['cate_estimates'].mean()

```python
from sklearn.cluster import KMeans
```

```python
kmeans = KMeans(n_clusters=2)
prediction = kmeans.fit_predict(data[common_causes].values)

```

```python
cl1 = prediction == 1
data1 = data[cl1].reset_index(drop = True)
data2 = data[~cl1].reset_index(drop = True)
```

```python
data1.describe()
```

      f0
      f1
      f2
      f3
      f4
      S
      Yf
      Ycf
      Y
      propensity_score
      strata
      dbar
      d_y
      dbar_y

      count
      531616.000000
      531616.000000
      531616.000000
      531616.000000
      531616.000000
      531616.0
      531616.0
      531616.0
      531616.0
      531616.000000
      531616.000000
      531616.000000
      531616.000000
      531616.000000

      mean
      -0.287588
      -0.348828
      -0.371957
      -0.316373
      -0.281822
      0.0
      1.0
      1.0
      1.0
      0.000408
      13.412177
      0.999426
      0.000574
      0.999426

      std
      0.950336
      0.926534
      0.915637
      0.939860
      0.950877
      0.0
      0.0
      0.0
      0.0
      0.000171
      7.890180
      0.023946
      0.023946
      0.023946

      min
      -4.895454
      -4.852118
      -4.802920
      -4.820940
      -5.031764
      0.0
      1.0
      1.0
      1.0
      0.000013
      0.000000
      0.000000
      0.000000
      0.000000

      25%
      -0.927748
      -0.971112
      -0.986030
      -0.947566
      -0.922035
      0.0
      1.0
      1.0
      1.0
      0.000272
      7.000000
      1.000000
      0.000000
      1.000000

      50%
      -0.284145
      -0.343608
      -0.365165
      -0.313446
      -0.279516
      0.0
      1.0
      1.0
      1.0
      0.000401
      13.000000
      1.000000
      0.000000
      1.000000

      75%
      0.355717
      0.277973
      0.249408
      0.319344
      0.361252
      0.0
      1.0
      1.0
      1.0
      0.000539
      20.000000
      1.000000
      0.000000
      1.000000

      max
      4.065940
      3.783802
      3.907693
      3.831790
      4.116536
      0.0
      1.0
      1.0
      1.0
      0.001065
      36.000000
      1.000000
      1.000000
      1.000000

```python
data1.Tr.sum(), ((data1.Tr == 1) & (data1.S==1)).sum()
```

    (305, 0)

```python
data2.Tr.sum(), ((data2.Tr == 1) & (data2.S==1)).sum()
```

    (695, 465)

#### Modeling of 2 clusters

```python
data_dict1 = {'df':data1,
             'treatment_name': 'Tr',
             'outcome_name': 'Y',
             'common_causes_names': common_causes,
             'time_val': None,
             'instrument_names': None,
             'dot_graph': None,
             'gml_graph': None,
             'ate': None}

data_dict2 = {'df':data2,
             'treatment_name': 'Tr',
             'outcome_name': 'Y',
             'common_causes_names': common_causes,
             'time_val': None,
             'instrument_names': None,
             'dot_graph': None,
             'gml_graph': None,
             'ate': None}

model1 = CausalModel(
        data=data_dict1['df'],
        treatment=data_dict1["treatment_name"],
        outcome=data_dict1["outcome_name"],
        common_causes=data_dict1["common_causes_names"],
        )

model2 = CausalModel(
        data=data_dict2['df'],
        treatment=data_dict2["treatment_name"],
        outcome=data_dict2["outcome_name"],
        common_causes=data_dict2["common_causes_names"],
        )

identified_estimand1 = model1.identify_effect(proceed_when_unidentifiable=True)
identified_estimand2 = model2.identify_effect(proceed_when_unidentifiable=True)

```

    WARNING:dowhy.causal_model:Causal Graph not provided. DoWhy will construct a graph based on data inputs.
    WARNING:dowhy.causal_model:Causal Graph not provided. DoWhy will construct a graph based on data inputs.
    WARNING:dowhy.causal_identifier:If this is observed data (not from a randomized experiment), there might always be missing confounders. Causal effect cannot be identified perfectly.
    WARNING:dowhy.causal_identifier:If this is observed data (not from a randomized experiment), there might always be missing confounders. Causal effect cannot be identified perfectly.

```python
estimate1 = model1.estimate_effect(identified_estimand1,
                                 method_name="backdoor.propensity_score_stratification",
                                target_units='ate',
                                test_significance = None,
                                evaluate_effect_strength = False,
 confidence_intervals=False)
print(estimate1)
print("Causal Estimate is " + str(estimate1.value))
```

    *** Causal Estimate ***

    ## Identified estimand
    Estimand type: nonparametric-ate
    ### Estimand : 1
    Estimand name: backdoor
    Estimand expression:
      d
    ─────(Expectation(Y|f3,f4,f0,f1,f2))
    d[Tr]
    Estimand assumption 1, Unconfoundedness: If U→{Tr} and U→Y then P(Y|Tr,f3,f4,f0,f1,f2,U) = P(Y|Tr,f3,f4,f0,f1,f2)
    ### Estimand : 2
    Estimand name: iv
    No such variable found!

    ## Realized estimand
    b: Y~Tr+f3+f4+f0+f1+f2
    Target units: ate

    ## Estimate
    Mean value: 0.0

    Causal Estimate is 0.0

```python
estimate2 = model2.estimate_effect(identified_estimand2,
                                 method_name="backdoor.propensity_score_stratification",
                                target_units='ate',
                                test_significance = None,
                                evaluate_effect_strength = False,
 confidence_intervals=False)
print(estimate2)
print("Causal Estimate is " + str(estimate2.value))
```

    *** Causal Estimate ***

    ## Identified estimand
    Estimand type: nonparametric-ate
    ### Estimand : 1
    Estimand name: backdoor
    Estimand expression:
      d
    ─────(Expectation(Y|f3,f4,f0,f1,f2))
    d[Tr]
    Estimand assumption 1, Unconfoundedness: If U→{Tr} and U→Y then P(Y|Tr,f3,f4,f0,f1,f2,U) = P(Y|Tr,f3,f4,f0,f1,f2)
    ### Estimand : 2
    Estimand name: iv
    No such variable found!

    ## Realized estimand
    b: Y~Tr+f3+f4+f0+f1+f2
    Target units: ate

    ## Estimate
    Mean value: 0.6641850953003251

    Causal Estimate is 0.6641850953003251

```python

```