<jupyter_start><jupyter_text>LocalizationUse "localization" to learn a Cahn-Hilliard model.  Learn a Cahn-HilliardSquare domain periodic boundary conditions $$ \dot{\phi} = \nabla^2 \left( \phi^3 - \phi \right) - \gamma \nabla^4 \phi $$ What are we trying to do?Create a mapping from $t_0$ to $t_{10}$ without doing all the steps. We want to do the following. $$ \phi[s](t=t_0) \rightarrow \phi[s](t=t_{10})$$  LocalizationUse regression for each local state.  Create Samples<jupyter_code>%matplotlib inline

import pymks
import matplotlib.pyplot as plt
import numpy as np
from pymks.datasets import make_cahn_hilliard
n_steps = 10
size = (151, 151)
X, y = make_cahn_hilliard(n_samples=10, size=size, dt=1., n_steps=n_steps)
print(X.shape)
print(y.shape)
# NBVAL_IGNORE_OUTPUT

plt.imshow(X[0])
plt.colorbar()
# NBVAL_IGNORE_OUTPUT

plt.imshow(y[0])
plt.colorbar()<jupyter_output><empty_output><jupyter_text>Parallel<jupyter_code>from dask import compute, delayed
import dask.multiprocessing

def make_data(seed):
    np.random.seed(seed)
    return make_cahn_hilliard(n_samples=10, size=size, dt=1., n_steps=n_steps)

funcs = [delayed(make_data)(seed) for seed in range(30)]
    
out = compute(*funcs, get=dask.multiprocessing.get)

np.array(out).shape
X = np.array(out)[:, 0].reshape((300,) + size)
y = np.array(out)[:, 1].reshape((300,) + size)<jupyter_output><empty_output><jupyter_text>Learning<jupyter_code>from pymks import MKSLocalizationModel
from pymks.bases import PrimitiveBasis

basis = PrimitiveBasis(n_states=5, domain=[-1, 1])
model = MKSLocalizationModel(basis=basis)
model.fit(X[:-1], y[:-1])
y_pred = model.predict(X[-1:])
# NBVAL_IGNORE_OUTPUT
plt.imshow(y_pred[0])
plt.colorbar()
# NBVAL_IGNORE_OUTPUT
plt.imshow(y[-1])
plt.colorbar()<jupyter_output><empty_output><jupyter_text>Train Test Split<jupyter_code>from sklearn.model_selection import train_test_split
from sklearn import metrics

X_train, X_test, y_train, y_test = train_test_split(X, y)
basis = PrimitiveBasis(n_states=5, domain=[-1, 1])
model = MKSLocalizationModel(basis=basis)
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
# NBVAL_IGNORE_OUTPUT

metrics.mean_squared_error(y_pred.flatten(), y_test.flatten())
# NBVAL_IGNORE_OUTPUT

print(y_pred[0][0][:10])
print(y_test[0][0][:10])
y_pred.shape
# NBVAL_IGNORE_OUTPUT

plt.imshow(y_pred[0])
plt.colorbar()
# NBVAL_IGNORE_OUTPUT

plt.imshow(y_test[0])
plt.colorbar()<jupyter_output><empty_output><jupyter_text>Scale Up<jupyter_code>X_big, y_big = make_cahn_hilliard(n_samples=1, size=(1000, 1000), dt=1., n_steps=n_steps)
# NBVAL_IGNORE_OUTPUT

plt.imshow(y_big[0])
plt.colorbar()
basis = PrimitiveBasis(n_states=5, domain=[-1, 1])
model = MKSLocalizationModel(basis=basis)
model.fit(X, y)
model.resize_coeff(y_big[0].shape)
y_big_pred = model.predict(X_big)
# NBVAL_IGNORE_OUTPUT

plt.imshow(y_big_pred[0])
plt.colorbar()
# NBVAL_IGNORE_OUTPUT

metrics.mean_squared_error(y_big_pred.flatten(), y_big.flatten())
# NBVAL_IGNORE_OUTPUT

%timeit make_cahn_hilliard(n_samples=1, size=(1000, 1000), dt=1., n_steps=n_steps)
# NBVAL_IGNORE_OUTPUT

%timeit model.predict(X_big)<jupyter_output>218 ms ± 2.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>Multiple Steps<jupyter_code>X2, y2 = make_cahn_hilliard(n_samples=1, size=size, dt=1., n_steps=2 * n_steps)

basis = PrimitiveBasis(n_states=10, domain=[-1, 1])
model = MKSLocalizationModel(basis=basis)
model.fit(X, y)
tmp = model.predict(X2)
y2_pred = model.predict(tmp)
# NBVAL_IGNORE_OUTPUT

plt.imshow(y2[0])
plt.colorbar()
# NBVAL_IGNORE_OUTPUT

plt.imshow(y2_pred[0])
plt.colorbar()
# NBVAL_IGNORE_OUTPUT

metrics.mean_squared_error(y2_pred.flatten(), y2.flatten())<jupyter_output><empty_output><jupyter_text>Cross Validation<jupyter_code>from pymks.bases import LegendreBasis
from sklearn.model_selection import GridSearchCV
from dask_searchcv import GridSearchCV
from sklearn import metrics
mse = metrics.mean_squared_error

prim_basis = PrimitiveBasis(2, [-1, 1])
leg_basis = LegendreBasis(2, [-1, 1])

params_to_tune = {'n_states': [2, 3, 5, 8, 13],
                  'basis': [prim_basis, leg_basis]}
model = MKSLocalizationModel(prim_basis)
score_func = metrics.make_scorer(lambda x, y: -mse(x.flatten(), y.flatten()))
gscv = GridSearchCV(model, params_to_tune, cv=5, scoring=score_func, scheduler="multiprocessing", n_jobs=8)
?GridSearchCV
# NBVAL_SKIP

gscv.fit(X_train, y_train)
# NBVAL_SKIP

gscv.best_estimator_
# NBVAL_SKIP

gscv.score(X_test, y_test)
# NBVAL_SKIP

gscv.cv_results_<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13562307353588726, 'fraction_numerical': 0.02664024658740643, 'mean_word_length': 7.170863309352518, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 28, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '628995', 'license': 'MIT', 'n_tokens_mistral': 1988, 'n_tokens_neox': 1840, 'n_words': 327, 'path': 'notebooks/part3-localization.ipynb', 'repo_name': 'auag92/mks_tutorial', 'size': '4542'}