import numpy as np
import torch
import sys

sys.path.append("../../")

import heat as ht
from matplotlib import pyplot as plt
from sklearn import datasets
import heat.ml.regression.lasso as lasso
import plotfkt

# read scikit diabetes data set
diabetes = datasets.load_diabetes()

# load diabetes dataset from hdf5 file
X = ht.load_hdf5("../../heat/datasets/data/diabetes.h5", dataset="x", split=0)
y = ht.load_hdf5("../../heat/datasets/data/diabetes.h5", dataset="y", split=0)

# normalize dataset #DoTO this goes into the lasso fit routine soon as issue #106 is solved
X = X / ht.sqrt((ht.mean(X ** 2, axis=0)))

# HeAT lasso instance
estimator = lasso.HeatLasso(max_iter=100)

# List  lasso model parameters
theta_list = list()

# Range of lambda values
lamda = np.logspace(0, 4, 10) / 10

# compute the lasso path
for l in lamda:
    estimator.lam = l
    estimator.fit(X, y)
    theta_list.append(estimator.theta.numpy().flatten())

# Stack estimated model parameters into one numpy array
theta_lasso = np.stack(theta_list).T

# Stack into numpy array
theta_lasso = np.stack(theta_list).T[1:, :]


# plot lasso paths
plt.subplot(3, 1, 1)
plotfkt.plot_lasso_path(
    lamda, theta_lasso, diabetes.feature_names, title="Lasso Paths - HeAT implementation"
)

if X.is_distributed():
    distributed = X.comm.rank
else:
    distributed = False

# Now the same stuff in numpy
X = diabetes.data.astype("float32")
y = diabetes.target.astype("float32")

m, _ = X.shape
X = np.concatenate((np.ones((m, 1)).astype("float32"), X), axis=1)

# normalize dataset
X = X / np.sqrt((np.mean(X ** 2, axis=0)))

# Numpy lasso instance
estimator = lasso.NumpyLasso(max_iter=100)

# List  lasso model parameters
theta_list = list()

# Range of lambda values
lamda = np.logspace(0, 4, 10) / 10

# compute the lasso path
for l in lamda:
    estimator.lam = l
    estimator.fit(X, y)
    theta_list.append(estimator.theta.flatten())

# Stack estimated model parameters into one numpy array
theta_lasso = np.stack(theta_list).T

# Stack into numpy array
theta_lasso = np.stack(theta_list).T[1:, :]

# plot lasso paths
plt.subplot(3, 1, 2)
plotfkt.plot_lasso_path(
    lamda, theta_lasso, diabetes.feature_names, title="Lasso Paths - Numpy implementation"
)

# Now the same stuff again in PyTorch
X = torch.tensor(X)
y = torch.tensor(y)

# HeAT lasso instance
estimator = lasso.PytorchLasso(max_iter=100)

# List lasso model parameters
theta_list = list()

# Range of lambda values
lamda = np.logspace(0, 4, 10) / 10

# compute the lasso path
for l in lamda:
    estimator.lam = l
    estimator.fit(X, y)
    theta_list.append(estimator.theta.numpy().flatten())

# Stack estimated model parameters into one numpy array
theta_lasso = np.stack(theta_list).T

# Stack into numpy array
theta_lasso = np.stack(theta_list).T[1:, :]

# plot lasso paths
plt.subplot(3, 1, 3)
plotfkt.plot_lasso_path(
    lamda, theta_lasso, diabetes.feature_names, title="Lasso Paths - PyTorch implementation"
)

# plot only with first rank
if distributed is False:
    plt.show()
elif distributed == 0:
    plt.show()

examples/lasso/demo.py
import numpy as np
import torch
import sys

sys.path.append("../../")

import heat as ht
from matplotlib import pyplot as plt
from sklearn import datasets
import heat.core.regression.lasso as lasso
import plotfkt

# read scikit diabetes data set
diabetes = datasets.load_diabetes()

# load diabetes dataset from hdf5 file
X = ht.load_hdf5("../../heat/datasets/data/diabetes.h5", dataset="x", split=0)
y = ht.load_hdf5("../../heat/datasets/data/diabetes.h5", dataset="y", split=0)

# normalize dataset #DoTO this goes into the lasso fit routine soon as issue #106 is solved
X = X / ht.sqrt((ht.mean(X ** 2, axis=0)))

# HeAT lasso instance
estimator = lasso.HeatLasso(max_iter=100)

# List  lasso model parameters
theta_list = list()

# Range of lambda values
lamda = np.logspace(0, 4, 10) / 10

# compute the lasso path
for l in lamda:
    estimator.lam = l
    estimator.fit(X, y)
    theta_list.append(estimator.theta.numpy().flatten())

# Stack estimated model parameters into one numpy array
theta_lasso = np.stack(theta_list).T

# Stack into numpy array
theta_lasso = np.stack(theta_list).T[1:, :]


# plot lasso paths
plt.subplot(3, 1, 1)
plotfkt.plot_lasso_path(
    lamda, theta_lasso, diabetes.feature_names, title="Lasso Paths - HeAT implementation"
)

if X.is_distributed():
    distributed = X.comm.rank
else:
    distributed = False

# Now the same stuff in numpy
X = diabetes.data.astype("float32")
y = diabetes.target.astype("float32")

m, _ = X.shape
X = np.concatenate((np.ones((m, 1)).astype("float32"), X), axis=1)

# normalize dataset
X = X / np.sqrt((np.mean(X ** 2, axis=0)))

# Numpy lasso instance
estimator = lasso.NumpyLasso(max_iter=100)

# List  lasso model parameters
theta_list = list()

# Range of lambda values
lamda = np.logspace(0, 4, 10) / 10

# compute the lasso path
for l in lamda:
    estimator.lam = l
    estimator.fit(X, y)
    theta_list.append(estimator.theta.flatten())

# Stack estimated model parameters into one numpy array
theta_lasso = np.stack(theta_list).T

# Stack into numpy array
theta_lasso = np.stack(theta_list).T[1:, :]

# plot lasso paths
plt.subplot(3, 1, 2)
plotfkt.plot_lasso_path(
    lamda, theta_lasso, diabetes.feature_names, title="Lasso Paths - Numpy implementation"
)

# Now the same stuff again in PyTorch
X = torch.tensor(X)
y = torch.tensor(y)

# HeAT lasso instance
estimator = lasso.PytorchLasso(max_iter=100)

# List lasso model parameters
theta_list = list()

# Range of lambda values
lamda = np.logspace(0, 4, 10) / 10

# compute the lasso path
for l in lamda:
    estimator.lam = l
    estimator.fit(X, y)
    theta_list.append(estimator.theta.numpy().flatten())

# Stack estimated model parameters into one numpy array
theta_lasso = np.stack(theta_list).T

# Stack into numpy array
theta_lasso = np.stack(theta_list).T[1:, :]

# plot lasso paths
plt.subplot(3, 1, 3)
plotfkt.plot_lasso_path(
    lamda, theta_lasso, diabetes.feature_names, title="Lasso Paths - PyTorch implementation"
)

# plot only with first rank
if distributed is False:
    plt.show()
elif distributed == 0:
    plt.show()

examples/lasso/demo.py
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000000..31ef645a82
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,5 @@
+# Format of .gitattributes is
+# <pattern> <attr1> <attr2> ...
+#
+
+CHANGELOG.md  merge=union
diff --git a/CHANGELOG.md b/CHANGELOG.md
index dfbd8a1015..f2c0c3b11c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+# Pending Additions
+
+- [#454](https://github.com/helmholtz-analytics/heat/issues/454) Update lasso example
+
 # v0.2.2
 
 This version adds support for PyTorch 1.4.0. There are also several minor feature improvements and bug fixes listed below.
diff --git a/examples/lasso/demo.py b/examples/lasso/demo.py
index 7cf89aae6f..b78cb6cd30 100644
--- a/examples/lasso/demo.py
+++ b/examples/lasso/demo.py
@@ -7,7 +7,7 @@
 import heat as ht
 from matplotlib import pyplot as plt
 from sklearn import datasets
-import heat.ml.regression.lasso as lasso
+import heat.core.regression.lasso as lasso
 import plotfkt
 
 # read scikit diabetes data set

Update lasso example
**Description**
The lasso example references heat.ml.regression.lasso

**To Reproduce**
Steps to reproduce the behavior:
1. Which module/class/function is affected?
  examples/lasso/demo.py
2. What are the circumstances under which the bug appears?
  run demo.py
3. What is the exact error-message/errorous behavious?
  `Traceback (most recent call last):
  File "demo.py", line 10, in <module>
    import heat.ml.regression.lasso as lasso
ModuleNotFoundError: No module named 'heat.ml'`

**Expected behavior**
Should import heat.core.regression.lasso

**Version Info**
 Which version are you using?
0.2.1

Update lasso example
**Description**
The lasso example references heat.ml.regression.lasso

**To Reproduce**
Steps to reproduce the behavior:
1. Which module/class/function is affected?
  examples/lasso/demo.py
2. What are the circumstances under which the bug appears?
  run demo.py
3. What is the exact error-message/errorous behavious?
  `Traceback (most recent call last):
  File "demo.py", line 10, in <module>
    import heat.ml.regression.lasso as lasso
ModuleNotFoundError: No module named 'heat.ml'`

**Expected behavior**
Should import heat.core.regression.lasso

**Version Info**
 Which version are you using?
0.2.1

