import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import segyio
import matplotlib.cm as cm
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsRegressor
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.callbacks import EarlyStopping
from mpl_toolkits.mplot3d import Axes3D  

# -------------- 1. Загрузка и предобработка данных --------------

segy_path = 'Cube.sgy'              
horizon_b_path = 'B'                
thickness_path = 'dh_dataset.txt'   

thickness_df = pd.read_csv(thickness_path, sep=r'\s+', header=None, names=['inline', 'crossline', 'thickness'])
horizon_b_df = pd.read_csv(horizon_b_path, sep=r'\s+', header=None, names=['x', 'y', 'inline', 'crossline', 'time_ms', 'unknown'])

merged_df = thickness_df.merge(horizon_b_df, on=['inline', 'crossline'], how='left')

# -------------- 2. Загрузка и обработка SEGY данных --------------

with segyio.open(segy_path, "r", ignore_geometry=True) as seg:
    ilines = seg.attributes(segyio.TraceField.INLINE_3D)[:]        
    xlines = seg.attributes(segyio.TraceField.CROSSLINE_3D)[:]     
    sample_rate_us = segyio.tools.dt(seg)
    sample_rate_ms = sample_rate_us / 1000

    traces = np.array([tr for tr in seg.trace])

il_xl_to_trace_idx = {(il, xl): idx for idx, (il, xl) in enumerate(zip(ilines, xlines))}

selected_traces = []
selected_coords = []

for _, row in merged_df.iterrows():
    il = row['inline']
    xl = row['crossline']
    time_ms = row['time_ms']
    trace_idx = il_xl_to_trace_idx.get((il, xl))
    trace = traces[trace_idx]

    window_size_samples = 30
    half_window = window_size_samples // 2  
    center_sample = int(time_ms / sample_rate_ms)
    start_sample = center_sample - half_window
    end_sample = center_sample + half_window
    trace_window = trace[start_sample:end_sample]

    rms = np.sqrt(np.mean(trace_window**2))
    trace_normalized = trace_window / rms

    selected_traces.append(trace_normalized)
    selected_coords.append((il, xl))

selected_traces = np.array(selected_traces)
selected_coords = np.array(selected_coords)

# -------------- 3. Сокращение размерности --------------

pca = PCA(n_components=3)
pca_components = pca.fit_transform(selected_traces)

input_dim = selected_traces.shape[1]
encoding_dim = 3

input_layer = Input(shape=(input_dim,))
encoded = Dense(16, activation='relu')(input_layer)
encoded = Dense(encoding_dim, activation='relu')(encoded)

decoded = Dense(16, activation='relu')(encoded)
decoded = Dense(input_dim, activation='linear')(decoded)

autoencoder = Model(inputs=input_layer, outputs=decoded)
encoder_model = Model(inputs=input_layer, outputs=encoded)

autoencoder.compile(optimizer='adam', loss='mse')

X_train_ae, X_val_ae = train_test_split(selected_traces, test_size=0.2, random_state=42)

early_stop = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)

history = autoencoder.fit(
    X_train_ae, X_train_ae,
    epochs=50,
    batch_size=32,
    shuffle=True,
    validation_data=(X_val_ae, X_val_ae),
    callbacks=[early_stop],
    verbose=1
)

plt.figure(figsize=(10, 6))
plt.plot(history.history['loss'], label='Train Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.title('Обучение Автокодировщика')
plt.xlabel('Эпоха')
plt.ylabel('MSE Loss')
plt.legend()

autoencoder_features = encoder_model.predict(selected_traces)

dataset_pca = pd.DataFrame(pca_components, columns=['PC1', 'PC2', 'PC3'])
dataset_pca['thickness'] = merged_df['thickness'].values
dataset_pca['inline'] = selected_coords[:, 0]
dataset_pca['crossline'] = selected_coords[:, 1]

dataset_autoencoder = pd.DataFrame(autoencoder_features, columns=['AE1', 'AE2', 'AE3'])
dataset_autoencoder['thickness'] = merged_df['thickness'].values
dataset_autoencoder['inline'] = selected_coords[:, 0]
dataset_autoencoder['crossline'] = selected_coords[:, 1]

# -------------- 4. Обучение и прогнозирование KNN для PCA и автокодировщика --------------

X_pca = dataset_pca[['PC1', 'PC2', 'PC3']].values
y_pca = dataset_pca['thickness'].values
X_train_pca, X_test_pca, y_train_pca, y_test_pca = train_test_split(X_pca, y_pca, test_size=0.2, random_state=42)
knn_pca = KNeighborsRegressor(n_neighbors=5)
knn_pca.fit(X_train_pca, y_train_pca)
y_pred_pca = knn_pca.predict(X_test_pca)
mse_pca = mean_squared_error(y_test_pca, y_pred_pca)
print(f"MSE для PCA: {mse_pca:.4f}")

X_ae = dataset_autoencoder[['AE1', 'AE2', 'AE3']].values
y_ae = dataset_autoencoder['thickness'].values
X_train_ae, X_test_ae, y_train_ae, y_test_ae = train_test_split(X_ae, y_ae, test_size=0.2, random_state=42)
knn_ae = KNeighborsRegressor(n_neighbors=5)
knn_ae.fit(X_train_ae, y_train_ae)
y_pred_ae = knn_ae.predict(X_test_ae)
mse_ae = mean_squared_error(y_test_ae, y_pred_ae)
print(f"MSE для автокодировщика: {mse_ae:.4f}")

# -------------- 5. Визуализация результатов --------------

fig_pca = plt.figure(figsize=(10, 8))
ax_pca = fig_pca.add_subplot(111, projection='3d')
scatter_pca = ax_pca.scatter(
    dataset_pca['PC1'], 
    dataset_pca['PC2'], 
    dataset_pca['PC3'], 
    c=dataset_pca['thickness'], 
    cmap='turbo', 
)
ax_pca.set_title('Кроссплот PCA компонентов с толщиной горизонта')
ax_pca.set_xlabel('PC1')
ax_pca.set_ylabel('PC2')
ax_pca.set_zlabel('PC3')
fig_pca.colorbar(scatter_pca, ax=ax_pca, pad=0.1)

fig_ae = plt.figure(figsize=(10, 8))
ax_ae = fig_ae.add_subplot(111, projection='3d')
scatter_ae = ax_ae.scatter(
    dataset_autoencoder['AE1'], 
    dataset_autoencoder['AE2'], 
    dataset_autoencoder['AE3'], 
    c=dataset_autoencoder['thickness'], 
    cmap='turbo', 
)
ax_ae.set_title('Кроссплот Компонентов Автокодировщика с толщиной горизонта')
ax_ae.set_xlabel('AE1')
ax_ae.set_ylabel('AE2')
ax_ae.set_zlabel('AE3')
fig_ae.colorbar(scatter_ae, ax=ax_ae, pad=0.1)

plt.figure(figsize=(12, 6))
scatter_map_pca = plt.scatter(
    dataset_pca['inline'], 
    dataset_pca['crossline'], 
    c=knn_pca.predict(dataset_pca[['PC1', 'PC2', 'PC3']].values), 
    cmap='turbo'
)
plt.colorbar(scatter_map_pca, label='Прогнозная толщина (м)')
plt.title('Карта прогнозных толщин (PCA + KNN) - Весь датасет')
plt.xlabel('Inline')
plt.ylabel('Crossline')

plt.figure(figsize=(12, 6))
scatter_map_ae = plt.scatter(
    dataset_autoencoder['inline'], 
    dataset_autoencoder['crossline'], 
    c=knn_ae.predict(dataset_autoencoder[['AE1', 'AE2', 'AE3']].values), 
    cmap='turbo'
)
plt.colorbar(scatter_map_ae, label='Прогнозная толщина (м)')
plt.title('Карта прогнозных толщин (Autoencoder + KNN) - Весь датасет')
plt.xlabel('Inline')
plt.ylabel('Crossline')

plt.show()

Добрый вечер! Основной момент: вычисление главных компонент и обучение автокодировщика должно проводиться не по 1000 трассам dh_dataset.txt, а по всем трассам (>1000000) B. На кроссплотах как будто мало точек из-за большего количества повторов значений для PC1,PC2,PC3.  Если сделаете расчет по всем трассам, то должно стать лучше. Ну и карты следует строить не по 1000 значениям, что были в обучении, а по всем трассам.