spiralsLFBGS.py

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.colors as mcolors
import numpy as np
import matplotlib.pyplot as plt

#parameters
R = 4   
k = 8
N = R
S = k * N
Nt = S * N

# inizialization
x0 = 0
y0 = 0
X = np.zeros((N, S))
Y = np.zeros((N, S))
Regions = np.zeros((N, S)) 

#coord
for s in range(1, S + 1):
    r = np.floor((s - 1) * R / S + 1)
    for n in range(1, N + 1):
        a0 = n / N * 2 * np.pi / k
        a = a0 + s * 2 * np.pi / S 
        b = a0
        
        
        X[n-1, s-1] = b * np.cos(a) + x0
        Y[n-1, s-1] = b * np.sin(a) + y0
        Regions[n-1, s-1] = r


Coordinates = np.vstack((X.flatten('F'), Y.flatten('F')))
Regions_flat = Regions.flatten('F')

# draw spiral
plt.figure(figsize=(6, 6), facecolor='white')
plt.axis('off')

colors = ['m', 'g', 'c', 'y']


for r in range(1, R + 1):
    I = (Regions_flat == r)
    plt.plot(Coordinates[0, I], Coordinates[1, I], 'o', 
             markeredgecolor='k', 
             markerfacecolor=colors[r-1], 
             markersize=6)

plt.axis('equal')
plt.title('Dane Treningowe - Spirale ')
plt.show()
# one hot encoding
Probabilities = np.zeros((R, Nt))
for n in range(Nt):
    
    Probabilities[int(Regions_flat[n]) - 1, n] = 1

X_tensor = torch.tensor(Coordinates.T, dtype=torch.float32)
Y_tensor = torch.tensor(Probabilities.T, dtype=torch.float32)

#model
model = nn.Sequential(
    nn.Linear(2, 16),
    nn.Tanh(),        
    nn.Linear(16, 8),
    nn.Tanh(),        
    nn.Linear(8, R),  
    nn.Tanh(),        
    nn.Linear(R, R),  
)

criterion = nn.MSELoss()


epochs = 200
min_mse = 1e-2


optimizer = optim.LBFGS(model.parameters(), lr=0.1, max_iter=20)
loss_history = []

#training
for epoch in range(epochs):
    
    def closure():
        optimizer.zero_grad()
        outputs = model(X_tensor)
        loss = criterion(outputs, Y_tensor)
        loss.backward()
        return loss

    optimizer.step(closure)
    

    with torch.no_grad():
        outputs = model(X_tensor)
        loss = criterion(outputs, Y_tensor)
        avg_loss = loss.item()
        
    loss_history.append(avg_loss)
    
    
    if (epoch + 1) % 10 == 0 or epoch == 0:
        print(f"epoch {epoch+1}/{epochs}, MSE: {avg_loss:.4f}")
        
    
    if avg_loss < min_mse:
        break

print("Trenowanie zakończone.")

# loss
plt.figure(figsize=(8, 5))
plt.plot(range(1, len(loss_history) + 1), loss_history, marker='+', linestyle='-', color='b', label='Training') 
plt.axhline(y=min_mse, color='g', linestyle='-', label='Minimum MSE (1e-2)')
plt.yscale('log')
plt.grid(True, which="both", ls="-", alpha=0.5)
plt.xlabel('Epoch')
plt.ylabel('Mean square error')
plt.title('Training algorithm: L-BFGS')
plt.legend()
plt.show()
#test
N2 = 300 


M = (1 + 1 / (2 * N)) / k * 2 * np.pi
V_grid = np.linspace(-M, M, N2)


X2, Y2 = np.meshgrid(V_grid + x0, V_grid + y0, indexing='ij')


Coordinates2 = np.vstack((X2.flatten('F'), Y2.flatten('F')))


X2_tensor = torch.tensor(Coordinates2.T, dtype=torch.float32)

model.eval() 
with torch.no_grad(): 
    Probabilities2 = model(X2_tensor) 
    
    _, Regions2_pred = torch.max(Probabilities2, 1) 


Regions2_pred = Regions2_pred.numpy() + 1


Regions2_grid = Regions2_pred.reshape((N2, N2), order='F')


cmap_background = mcolors.ListedColormap(['red', 'green', 'blue', 'yellow'])

plt.figure(figsize=(8, 8), facecolor='white')
plt.axis('off')


plt.pcolormesh(X2, Y2, Regions2_grid, cmap=cmap_background, alpha=0.25, shading='auto')


colors_dots = ['m', 'g', 'c', 'y']
for r in range(1, R + 1):
    I = (Regions_flat == r)
    plt.plot(Coordinates[0, I], Coordinates[1, I], 'o', 
             markeredgecolor='k', markerfacecolor=colors_dots[r-1], markersize=6)

plt.axis('equal')
plt.title('Test Sieci: Wyuczone Granice Decyzyjne (L-BFGS)')
plt.show()