{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"},{"sourceId":11422292,"sourceType":"datasetVersion","datasetId":7153477},{"sourceId":11432031,"sourceType":"datasetVersion","datasetId":7160133}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport csv\n\nfrom matplotlib import pyplot as plt\nimport seaborn as sns\nfrom plotly.offline import init_notebook_mode\nimport plotly.offline as py\nimport plotly.express as px\nimport plotly.graph_objects as go\nfrom plotly.subplots import make_subplots\nfrom IPython.display import Image\n\nfrom tqdm import tqdm\nfrom tqdm.auto import tqdm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pathlib import Path\nfrom torch.utils.data import Dataset, DataLoader, random_split\n\nfrom sklearn.model_selection import KFold, train_test_split\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        os.path.join(dirname, filename)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:08:38.008323Z","iopub.execute_input":"2025-09-30T09:08:38.008665Z","iopub.status.idle":"2025-09-30T09:10:34.380924Z","shell.execute_reply.started":"2025-09-30T09:08:38.008608Z","shell.execute_reply":"2025-09-30T09:10:34.380103Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# <span style=\"font-size:20pt\">A bit of theory from seismic exploration</span>","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_1.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.382191Z","iopub.execute_input":"2025-09-30T09:10:34.382712Z","iopub.status.idle":"2025-09-30T09:10:34.393443Z","shell.execute_reply.started":"2025-09-30T09:10:34.382688Z","shell.execute_reply":"2025-09-30T09:10:34.392482Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"1.1. The simplest models of media used in seismic exploration for\nthe dependence of velocity and depth V(h):\n- a) is a homogeneous half–space with constant velocity;\n- b) – horizontally layered medium (SCS);\n- c) - vertically heterogeneous medium with variable velocity and constant\nvelocity gradient;\n- d) is a vertically inhomogeneous medium with variable velocity and a positive\nincreasing velocity gradient;\n- e) is a vertically inhomogeneous medium with variable velocity and a positive\ndecreasing velocity gradient.\n- f) is a vertically inhomogeneous medium with variable velocity and negative\nan increasing gradient;\n\nWe work with 3 types of models (it's easy to guess which ones)","metadata":{}},{"cell_type":"markdown","source":"**Types of seismic exploration**\n\nAccording to the location, there are: onshore, offshore, underground and borehole\nseismic surveys (VSP). According to the purpose and tasks to be solved, the following are distinguished:\n\n• shallow seismic exploration:\n- when studying the geological structure of the upper part\nof the section;\n- in construction to determine the elastic modules of rocks;\n- when karsts and suffusion hazards are detected;\n- when determining the groundwater level;\n- when studying the properties of loose sediments on the shelf;\n- when searching for gas hydrates.\n\n• Oil and gas seismic exploration:\n- in the search and exploration of hydrocarbons;\n- when determining and adjusting the locations of production wells;\n- when determining porosity and oil saturation parameters;\n- during monitoring of developed fields.\n\n• Structural seismic exploration:\n- when studying the deep structure of the Earth;\n- when mapping the foundation;\n- when mapping regional faults\n\n","metadata":{}},{"cell_type":"markdown","source":"**Types of seismic velocities used in seismic exploration**\n\nThe true velocity of a seismic wave is the velocity in an infinitesimally small volume of a medium, i.e., where its heterogeneity is least affected. The average seismic velocity is defined as the ratio of the sum of the layer capacities to the sum of the wave travel times in each layer.:","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_2.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.395821Z","iopub.execute_input":"2025-09-30T09:10:34.396078Z","iopub.status.idle":"2025-09-30T09:10:34.410511Z","shell.execute_reply.started":"2025-09-30T09:10:34.396059Z","shell.execute_reply":"2025-09-30T09:10:34.409703Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The interval velocity is the velocity at a certain depth range.:","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_3.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.412365Z","iopub.execute_input":"2025-09-30T09:10:34.412615Z","iopub.status.idle":"2025-09-30T09:10:34.427926Z","shell.execute_reply.started":"2025-09-30T09:10:34.412594Z","shell.execute_reply":"2025-09-30T09:10:34.426968Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The reservoir velocity is rigidly tied to a certain layer:","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_4.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.428849Z","iopub.execute_input":"2025-09-30T09:10:34.429427Z","iopub.status.idle":"2025-09-30T09:10:34.443359Z","shell.execute_reply.started":"2025-09-30T09:10:34.429403Z","shell.execute_reply":"2025-09-30T09:10:34.442560Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Another type of velocity is apparent velocity. This is what we determine when working on land or offshore on the observation line. The apparent speed can be calculated using the formula: ","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_5.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.444206Z","iopub.execute_input":"2025-09-30T09:10:34.444432Z","iopub.status.idle":"2025-09-30T09:10:34.459425Z","shell.execute_reply.started":"2025-09-30T09:10:34.444416Z","shell.execute_reply":"2025-09-30T09:10:34.458540Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"where dx is the distance between the receivers. The apparent velocity, depending on the angle of approach of the beam to the surface, obeys Bendorf's law:","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_6.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.460420Z","iopub.execute_input":"2025-09-30T09:10:34.460755Z","iopub.status.idle":"2025-09-30T09:10:34.477093Z","shell.execute_reply.started":"2025-09-30T09:10:34.460729Z","shell.execute_reply":"2025-09-30T09:10:34.476208Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"where a is the angle between the ray and the surface normal. It follows from this relation that the apparent velocity varies from the value of the true velocity to infinity.\n\nIn seismic exploration, we can plot the dependence of the arrival time of a wave on the coordinate of the geophone registering this wave. Such a graph will be called the hodograph of a particular wave. The apparent speed can be easily determined by its hodograph using the first formula.","metadata":{}},{"cell_type":"markdown","source":"**Seismic waves and their kinematic characteristics**\n\nThe seismogram and its elements. By placing the reception points on the profile, we can record the arrival time of the seismic wave.\n\nHodograph is the dependence of the time of arrival of a seismic wave on the coordinates of the receiving (removal) points. Having built a hodograph, we can proceed to calculating velocities using mathematical equations, which we will discuss later. If we place a geophone at some point that will register ground fluctuations, then at the output we will get a field record - a seismic trace.","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_7.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.478014Z","iopub.execute_input":"2025-09-30T09:10:34.478773Z","iopub.status.idle":"2025-09-30T09:10:34.492985Z","shell.execute_reply.started":"2025-09-30T09:10:34.478750Z","shell.execute_reply":"2025-09-30T09:10:34.492165Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The seismic path is the dependence of the amplitude (In or mV) of seismic events on the time of their registration (s or ms). The fastest way to reach the seismic receiver is a longitudinal wave, the first entry of which is clearly visible on the seismic path (a sharp deviation from zero of the soil particles). After a while, the longitudinal wave will fade out and the next one will arrive – the transverse wave. Next comes a huge amplitude of surface waves, which are characterized by high amplitude, low frequency and low speed. Such\nThe seismic event reflected on the graph is a seismic track or seismic path. Obviously, each reception point will record its own seismic path. The pulse amplitudes will be determined by the reflection coefficients.\n\nIn real seismic exploration, not one reception point is used, but many. By placing a certain number of geophones and activating the signal, we will get our own seismic path for each geophone. The totality (ensemble) The seismic traces obtained for a single excitation point are called a seismogram. The seismogram shows the numbers of reception points along the abscissa axis, usually in the same step, and the recording time, which is set by a human, on the other axis.","metadata":{}},{"cell_type":"markdown","source":"We only work with direct and reflected waves, I'll tell you more about them.\n\nA straight wave can be detected if we consider the first layer from the surface to be homogeneous. Let's say we have a source and receivers that register our signal. A straight line will be a wave propagating along the surface of the earth at a certain speed. The hodograph of this wave will be a straight line. If the task is to determine the travel time of the signal to the receiver with a certain coordinate, then the dependence will have the form:","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_8.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.494027Z","iopub.execute_input":"2025-09-30T09:10:34.494333Z","iopub.status.idle":"2025-09-30T09:10:34.502615Z","shell.execute_reply.started":"2025-09-30T09:10:34.494286Z","shell.execute_reply":"2025-09-30T09:10:34.501695Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"By specifying the distance explicitly, we get the time value at each point. \n\nSo, the hodograph of a straight wave has the form of a linear function. A plus sign is placed when we go to the right of the source, a minus sign is placed to the left. In space, the hodograph of a straight wave will be a cone. The higher the velocity in the medium, the more the hodograph is pressed against the Oh axis. The longitudinal wave hodograph will always be more pressed against the axis than the transverse wave hodograph, due to the difference in speeds for them. The apparent speed in this case is equal to the true speed. The front of a straight wave is a cylinder.\n\nThe reflected wave has a complex hodograph equation (a parabola with branches up), I will not give it. A reflected wave is a wave that has reflected from the layer only 1 time (more than 1 time is a multiple reflected). Consider their hodographs (direct and reflected waves):","metadata":{}},{"cell_type":"code","source":"Image(\"/kaggle/input/seismic-images-theory/Print_9.png\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:34.504433Z","iopub.execute_input":"2025-09-30T09:10:34.504724Z","iopub.status.idle":"2025-09-30T09:10:34.519682Z","shell.execute_reply.started":"2025-09-30T09:10:34.504695Z","shell.execute_reply":"2025-09-30T09:10:34.518691Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# <span style=\"font-size:20pt\">Stage 1: Exploratory Data Analysis (EDA)</span>","metadata":{}},{"cell_type":"code","source":"def perform_eda(data_path):\n    example_files = {\n        'Curve-Vel (data-model)': {'seismic': '/kaggle/input/waveform-inversion/train_samples/CurveVel_A/data/data2.npy', 'velocity': '/kaggle/input/waveform-inversion/train_samples/CurveVel_A/model/model2.npy'},\n        'Curve-Fault (seis-vel)': {'seismic': '/kaggle/input/waveform-inversion/train_samples/CurveFault_A/seis4_1_0.npy', 'velocity': '/kaggle/input/waveform-inversion/train_samples/CurveFault_A/vel4_1_0.npy'},\n        'Style (data-model)': {'seismic': '/kaggle/input/waveform-inversion/train_samples/Style_A/data/data2.npy','velocity': '/kaggle/input/waveform-inversion/train_samples/Style_A/model/model2.npy'}}\n    \n    plt.figure(figsize=(15, 10))\n    for i, (family, files) in enumerate(example_files.items(), 1):\n        # Uploading data with the correct path\n        seis_data = np.load(data_path / files['seismic'], mmap_mode='r')[0]\n        vel_map = np.load(data_path / files['velocity'], mmap_mode='r')[0]\n        \n        # Correcting the velocity map shape\n        if vel_map.ndim == 3:  # If the dimension is (1, H, W)\n            vel_map = vel_map.squeeze(0)  # We delete the first dimension\n        \n        # Graph 1: Slices of seismic data\n        plt.subplot(3, 3, i)\n        plt.imshow(seis_data[0].T, aspect='auto', cmap='seismic')  # Simplified access to data\n        plt.title(f'{family} - Сейсмические данные\\n(Источник 0)')\n        plt.xlabel('Time steps')\n        plt.ylabel('Receivers')\n        \n        # График 2: Velocity Map\n        plt.subplot(3, 3, i+3)\n        plt.imshow(vel_map, cmap='jet', vmin=1500, vmax=6000)  # Added a range of speeds\n        plt.title(f'{family} - Velocity Map')\n        plt.colorbar(label='Velocity (м/с)')\n        \n        # Graph 3: Velocity distribution\n        plt.subplot(3, 3, i+6)\n        sns.histplot(vel_map.ravel(), bins=50, kde=True)\n        plt.title(f'{family} - Распределение скоростей')\n        plt.xlabel('Скорость')\n        plt.ylabel('Частота')\n        plt.xlim(1500, 6000)  # Range limitation\n    \n    plt.tight_layout()\n    plt.show()\n\ndata_path = Path('/kaggle/input/waveform-inversion/train_samples')\nperform_eda(data_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:46.743957Z","iopub.execute_input":"2025-09-30T09:10:46.744250Z","iopub.status.idle":"2025-09-30T09:10:49.185997Z","shell.execute_reply.started":"2025-09-30T09:10:46.744225Z","shell.execute_reply":"2025-09-30T09:10:49.184984Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def perform_advanced_eda(data_path):\n    # Analyze different data families with proper relative file structure\n    families = {\n        'FlatVel': {\n            'path': 'FlatVel_{version}/data/data*.npy',\n            'velocity': 'FlatVel_{version}/model/model*.npy'\n        },\n        'CurveVel': {\n            'path': 'CurveVel_{version}/data/data*.npy',\n            'velocity': 'CurveVel_{version}/model/model*.npy'\n        },\n        'Style': {\n            'path': 'Style_{version}/data/data*.npy',\n            'velocity': 'Style_{version}/model/model*.npy'\n        },\n        'FlatFault': {\n            'path': 'FlatFault_{version}/seis*.npy',\n            'velocity': 'FlatFault_{version}/vel*.npy'\n        },\n        'CurveFault': {\n            'path': 'CurveFault_{version}/seis*.npy',\n            'velocity': 'CurveFault_{version}/vel*.npy'\n        }\n    }\n\n    # Create sample comparison across families\n    compare_families = [('FlatVel', 'A'),\n                        ('CurveVel', 'A'),\n                        ('Style', 'A'),\n                        ('CurveFault', 'A')]\n    \n    fig = make_subplots(\n        rows=2, cols=2,\n        subplot_titles=[f\"{name}_{ver}\" for name, ver in compare_families],\n        specs=[[{'type': 'heatmap'}, {'type': 'heatmap'}],\n               [{'type': 'heatmap'}, {'type': 'heatmap'}]]\n    )\n    \n    for i, (family, version) in enumerate(compare_families):\n        try:\n            # Generate relative path patterns\n            seis_pattern = families[family]['path'].format(version=version)\n            vel_pattern = families[family]['velocity'].format(version=version)\n            \n            # Find files relative to data_path\n            seis_files = sorted(data_path.glob(seis_pattern))\n            vel_files = sorted(data_path.glob(vel_pattern))\n            \n            if not seis_files or not vel_files:\n                print(f\"Missing files for {family}_{version}\")\n                continue\n\n            # Load sample data using first matching file\n            seis_data = np.load(seis_files[0], mmap_mode='r')[0]\n            vel_map = np.load(vel_files[0], mmap_mode='r')[0].squeeze()\n\n            row = i//2 + 1\n            col = i%2 + 1\n            \n            # Add visualizations\n            #fig.add_trace(go.Heatmap(z=seis_data[0,:,:].T, colorscale='Viridis'), row=row, col=col)\n            fig.add_trace(go.Heatmap(z=vel_map, colorscale='Jet', showscale=False),row=row, col=col)\n\n        except Exception as e:\n            print(f\"Error processing {family}_{version}: {str(e)}\")\n            continue\n    \n    fig.update_layout(title_text='Cross-Family Data Comparison', height=1000, width=1400)\n    #fig.show(renderer='iframe')\n    py.iplot(fig)\n\n# Execute EDA with corrected paths\ninit_notebook_mode(connected=True)\nperform_advanced_eda(Path('/kaggle/input/waveform-inversion/train_samples'))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:51.399656Z","iopub.execute_input":"2025-09-30T09:10:51.400452Z","iopub.status.idle":"2025-09-30T09:10:53.124967Z","shell.execute_reply.started":"2025-09-30T09:10:51.400421Z","shell.execute_reply":"2025-09-30T09:10:53.124153Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# <span style=\"font-size:20pt\">Stage 2: Configuration and data preparation</span>","metadata":{}},{"cell_type":"code","source":"class Config:\n    output_shape = (70, 70)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:58.454617Z","iopub.execute_input":"2025-09-30T09:10:58.454935Z","iopub.status.idle":"2025-09-30T09:10:58.459315Z","shell.execute_reply.started":"2025-09-30T09:10:58.454915Z","shell.execute_reply":"2025-09-30T09:10:58.458464Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_data(data_path):\n    # Uploading data (seis/data)\n    # data_path = '/kaggle/input/waveform-inversion/train_samples'\n    # Uploading paths to data files\n    all_inputs = [f for f in Path(data_path).rglob('*.npy') if ('seis' in f.stem) or ('data' in f.stem)]\n    \n    def get_output_files(input_files):\n        # Generating paths to target files\n        return [Path(str(f).replace('seis', 'vel').replace('data', 'model')) for f in input_files]\n    \n    all_outputs = get_output_files(all_inputs)\n    assert all(f.exists() for f in all_outputs)\n    return all_inputs, all_outputs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:10:58.656147Z","iopub.execute_input":"2025-09-30T09:10:58.656453Z","iopub.status.idle":"2025-09-30T09:10:58.662053Z","shell.execute_reply.started":"2025-09-30T09:10:58.656429Z","shell.execute_reply":"2025-09-30T09:10:58.661147Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# <span style=\"font-size:20pt\">Stage 3: Defining the dataset and model</span>","metadata":{}},{"cell_type":"code","source":"class SeismicDataset(Dataset):\n    def __init__(self, input_files, output_files, n_examples = 500, augment = True):\n        \n        self.input_files = input_files    # List of input files\n        self.output_files = output_files  # List of files with target data\n        self.n_examples = n_examples      # Number of examples in each file\n        self.augment = augment            # The augmentation flag\n        \n        # Calculation of statistics for normalization\n        sample = np.load(input_files[0], mmap_mode='r')[0]\n        self.input_mean = sample.mean()   # The average value of the input data\n        self.input_std = sample.std()     # Standard deviation of the input data\n        \n        sample = np.load(output_files[0], mmap_mode='r')[0]\n        self.output_mean = sample.mean()  # The average value of the target data\n        self.output_std = sample.std()    # The standard deviation of the target data\n\n\n    def __len__(self):\n        # Total number of examples\n        return len(self.input_files) * self.n_examples\n\n    def __getitem__(self, idx):\n        \n        # Getting one example of data\n        file_idx = idx // self.n_examples  # File index\n        sample_idx = idx % self.n_examples # The index of the example in the file\n        \n        # Loading data from memory mapping\n        X = np.load(self.input_files[file_idx], mmap_mode='r')[sample_idx]  # Uploading seismic data\n        y = np.load(self.output_files[file_idx], mmap_mode='r')[sample_idx] # Loading the velocity map\n        \n        # Normalization of data\n        X = (X - self.input_mean) / self.input_std  # Normalization of input data\n        y = (y - self.output_mean) / self.output_std # Normalization of target data\n        \n        # Augmentation: horizontal reflection\n        if self.augment and np.random.rand() > 0.5:\n            X = np.flip(X, axis=-1).copy()  # Reflection on the last axis\n            y = np.flip(y, axis=-1).copy()\n        \n        # Conversion to PyTorch tensors\n        X = torch.FloatTensor(X).mean(dim=0, keepdim=True) # Averaging by sources\n        y = torch.FloatTensor(y)                           # Target Data Conversion\n        \n        return X, y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:11:00.934872Z","iopub.execute_input":"2025-09-30T09:11:00.935675Z","iopub.status.idle":"2025-09-30T09:11:00.944329Z","shell.execute_reply.started":"2025-09-30T09:11:00.935638Z","shell.execute_reply":"2025-09-30T09:11:00.943463Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# <span style=\"font-size:20pt\">Stage 4: U-Net Model Architecture</span>","metadata":{}},{"cell_type":"code","source":"class SeismicUNet(nn.Module):\n    def __init__(self):\n        super().__init__()\n        \n        # Encoder blocks with downsampling\n        self.enc1 = nn.Sequential(\n            nn.Conv2d(1, 64, 3, padding=1),  # Initial feature extraction\n            nn.BatchNorm2d(64),               # Batch normalization\n            nn.ReLU(),                        # Non-linear activation\n            nn.Conv2d(64, 64, 3, padding=1),  # Feature refinement\n            nn.BatchNorm2d(64),\n            nn.ReLU(),\n            nn.MaxPool2d(2)                   # Spatial downsampling\n        )\n        \n        self.enc2 = nn.Sequential(\n            nn.Conv2d(64, 128, 3, padding=1), # Deeper feature extraction\n            nn.BatchNorm2d(128),\n            nn.ReLU(),\n            nn.Conv2d(128, 128, 3, padding=1),\n            nn.BatchNorm2d(128),\n            nn.ReLU(),\n            nn.MaxPool2d(2)\n        )\n        \n        # Bottleneck with expanded features\n        self.bottleneck = nn.Sequential(\n            nn.Conv2d(128, 256, 3, padding=1),  # High-level features\n            nn.BatchNorm2d(256),\n            nn.ReLU(),\n            nn.Conv2d(256, 256, 3, padding=1),\n            nn.BatchNorm2d(256),\n            nn.ReLU(),\n            nn.ConvTranspose2d(256, 128, 2, stride=2)  # Upsampling\n        )\n        \n        # Decoder blocks with skip connections\n        self.dec2 = nn.Sequential(\n            nn.Conv2d(256, 128, 3, padding=1),  # Combine features\n            nn.BatchNorm2d(128),\n            nn.ReLU(),\n            nn.Conv2d(128, 128, 3, padding=1),\n            nn.BatchNorm2d(128),\n            nn.ReLU(),\n            nn.ConvTranspose2d(128, 64, 2, stride=2)\n        )\n        \n        self.dec1 = nn.Sequential(\n            nn.Conv2d(128, 64, 3, padding=1),  # Final feature combination\n            nn.BatchNorm2d(64),\n            nn.ReLU(),\n            nn.Conv2d(64, 64, 3, padding=1),\n            nn.BatchNorm2d(64),\n            nn.ReLU(),\n            nn.Conv2d(64, 1, 1)                # Output channel reduction\n        )\n        \n        # Final size adjustment\n        self.final_adjust = nn.Sequential(\n            nn.Upsample(size=Config.output_shape, mode='bilinear'),  # Ensure output size\n            nn.Conv2d(1, 1, 1)                                       # Final adjustment\n        )\n    \n    def forward(self, x):\n        # Forward pass through the network\n        \n        # Encoder path\n        e1 = self.enc1(x)  # First encoding stage\n        e2 = self.enc2(e1)  # Second encoding stage\n        \n        # Bottleneck processing\n        bn = self.bottleneck(e2)  # Process features at lowest resolution\n        \n        # Decoder path with skip connections\n        d2 = self.dec2(torch.cat([bn, F.interpolate(e2, bn.shape[2:], mode='bilinear')], 1))\n        d1 = self.dec1(torch.cat([d2, F.interpolate(e1, d2.shape[2:], mode='bilinear')], 1))\n        \n        # Final output adjustment\n        return self.final_adjust(d1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:11:04.968455Z","iopub.execute_input":"2025-09-30T09:11:04.968783Z","iopub.status.idle":"2025-09-30T09:11:04.981211Z","shell.execute_reply.started":"2025-09-30T09:11:04.968759Z","shell.execute_reply":"2025-09-30T09:11:04.980237Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CombinedLoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.mae = nn.L1Loss()\n        self.mse = nn.MSELoss()\n    \n    def gradient_loss(self, pred, target):\n        # Adding a dimension check\n        if pred.dim() == 3:\n            pred = pred.unsqueeze(1)  # Add the channel if it is missing\n        if target.dim() == 3:\n            target = target.unsqueeze(1)\n            \n        # Size alignment\n        if pred.size()[-2:] != target.size()[-2:]:\n            pred = F.interpolate(pred, size=target.size()[-2:], mode = 'bilinear')\n        \n        # Calculation of gradients with correct indexes\n        dy_pred = pred[:, :, 1:, :] - pred[:, :, :-1, :]\n        dx_pred = pred[:, :, :, 1:] - pred[:, :, :, :-1]\n        \n        dy_target = target[:, :, 1:, :] - target[:, :, :-1, :]\n        dx_target = target[:, :, :, 1:] - target[:, :, :, :-1]\n        \n        return self.mae(dy_pred, dy_target) + self.mae(dx_pred, dx_target)\n    \n    def forward(self, pred, target):\n        # Reduction to the same dimension\n        if pred.dim() != target.dim():\n            if pred.dim() == 4:\n                pred = pred.squeeze(1)\n            else:\n                pred = pred.unsqueeze(1)\n        \n        return (self.mae(pred, target) + 0.3 * self.mse(pred, target) + 0.2 * self.gradient_loss(pred, target))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:11:05.272535Z","iopub.execute_input":"2025-09-30T09:11:05.273467Z","iopub.status.idle":"2025-09-30T09:11:05.281673Z","shell.execute_reply.started":"2025-09-30T09:11:05.273439Z","shell.execute_reply":"2025-09-30T09:11:05.280833Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Config:\n    batch_size = 32\n    lr = 1e-4\n    epochs = 50\n    early_stop = 10             # Early-Stop\n    num_workers = 4             # Optional parameter (calculated on GPU)\n    input_shape = (5, 1000, 70) # [sources, time_steps, receivers]\n    output_shape = (70, 70)     # velocity map\n    plot_every = 5","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:11:09.262020Z","iopub.execute_input":"2025-09-30T09:11:09.262318Z","iopub.status.idle":"2025-09-30T09:11:09.266809Z","shell.execute_reply.started":"2025-09-30T09:11:09.262286Z","shell.execute_reply":"2025-09-30T09:11:09.266090Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# <span style=\"font-size:20pt\">Stage 5: Training Configuration and Execution</span>","metadata":{}},{"cell_type":"code","source":"# Initialize model and training components\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')                   # Hardware selection\nmodel = SeismicUNet().to(device)                                                        # Model instantiation\noptimizer = torch.optim.AdamW(model.parameters(), lr = Config.lr, weight_decay = 5e-5)  # Optimizer\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience = 3)  # LR scheduler\ncriterion = CombinedLoss()                                                              # Custom loss function\n#criterion = nn.L1Loss()                                                                  # None-Custom loss function","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T09:11:11.639300Z","iopub.execute_input":"2025-09-30T09:11:11.640140Z","iopub.status.idle":"2025-09-30T09:11:14.601811Z","shell.execute_reply.started":"2025-09-30T09:11:11.640111Z","shell.execute_reply":"2025-09-30T09:11:14.600933Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Uploading data\ndata_path = '/kaggle/input/waveform-inversion/train_samples'\nall_inputs, all_outputs = load_data(data_path)\n\n# Initializing lists to store loss history\ntrain_history = []\nval_history = []\n\n# Creating a shared dataset\nfull_dataset = SeismicDataset(all_inputs, all_outputs)\ninput_mean = full_dataset.input_mean\ninput_std = full_dataset.input_std\n\n# Creating a tensor from all the targets\nall_targets = torch.stack([torch.FloatTensor(np.load(f)) for f in all_outputs])\noutput_mean = all_targets.mean().item()\noutput_std = all_targets.std().item()\n\n\n# Dividing the data into training and validation ones\ntrain_size = int(0.8 * len(full_dataset))\nval_size = len(full_dataset) - train_size\ntrain_ds, val_ds = random_split(full_dataset, [train_size, val_size])\n\n# Applying augmentation to a training dataset\n# If the augmentation is controlled by a flag, create new datasets with the necessary parameters\ntrain_ds.dataset.augment = True  # Let's assume that it is possible to change the parameter\nval_ds.dataset.augment = False   # Disabling augmentation for validation\n\n# Creating loaders\ntrain_loader = DataLoader(train_ds, batch_size = Config.batch_size, shuffle = True)\nval_loader = DataLoader(val_ds, batch_size = Config.batch_size, shuffle = False)\n\nbest_loss = float('inf')\ncounter = 0\n\n\n# The learning cycle\nfor epoch in range(Config.epochs):\n    \n    model.train()\n    epoch_train_loss = 0.0\n    epoch_val_loss = 0.0\n    for X, y in tqdm(train_loader, desc = f'Epoch {epoch+1} [Train]'):\n        X, y = X.to(device), y.to(device)\n        optimizer.zero_grad()\n        outputs = model(X)\n        loss = criterion(outputs, y)\n        loss.backward()\n        optimizer.step()\n        epoch_train_loss += loss.item() * X.size(0)\n    \n    model.eval()\n    with torch.no_grad():\n        for X, y in tqdm(val_loader, desc = f'Epoch {epoch+1} [Val]'):\n            X, y = X.to(device), y.to(device)\n            outputs = model(X)\n            loss = criterion(outputs, y)\n            epoch_val_loss += loss.item() * X.size(0)\n    \n    train_loss = epoch_train_loss / len(train_loader.dataset)\n    val_loss = epoch_val_loss / len(val_loader.dataset)\n    \n    # Saving history and logging remain unchanged\n    train_history.append(train_loss)\n    val_history.append(val_loss)\n\n    scheduler.step(val_loss)\n    \n    if (epoch+1) % Config.plot_every == 0:\n        plt.figure(figsize=(10, 5))\n        plt.plot(train_history, label='Train Loss')\n        plt.plot(val_history, label='Validation Loss')\n        plt.title('Training History')\n        plt.xlabel('Epoch')\n        plt.ylabel('Loss')\n        plt.legend()\n        plt.grid(True)\n        plt.show()\n    \n    print(f'Epoch {epoch+1}/{Config.epochs}')\n    print(f'Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}')\n    \n    if val_loss < best_loss:\n        best_loss = val_loss\n        torch.save(model.state_dict(), 'best_model_2.pth')\n        counter = 0\n    else:\n        counter += 1\n        if counter >= Config.early_stop:\n            print('Early stopping triggered')\n            break\n\n# Final visualization\nplt.figure(figsize=(12, 6))\nplt.plot(train_history, 'o-', label='Train Loss')\nplt.plot(val_history, 'o-', label='Validation Loss')\nplt.title('Training and Validation Loss History')\nplt.xlabel('Epoch')\nplt.ylabel('Loss')\nplt.legend()\nplt.grid(True)\nplt.close()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-17T07:59:07.432599Z","iopub.execute_input":"2025-04-17T07:59:07.433450Z","iopub.status.idle":"2025-04-17T07:59:08.800343Z","shell.execute_reply.started":"2025-04-17T07:59:07.433422Z","shell.execute_reply":"2025-04-17T07:59:08.799604Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# <span style=\"font-size:20pt\">Stage 6: Prediction and Submission Generation</span>","metadata":{}},{"cell_type":"code","source":"class TestDataset(Dataset):\n    # Dataset for loading and processing test data\n    def __init__(self, files):\n        self.files = files  # List of test files\n        # Use training statistics for normalization\n        self.input_mean = input_mean\n        self.input_std = input_std\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, idx):\n        # Load and process test sample\n        data = np.load(self.files[idx], mmap_mode = 'r')\n        data = (data - self.input_mean) / self.input_std  # Normalize\n        data = torch.FloatTensor(data.mean(axis=0, keepdims=True))  # Average sources\n        return data, self.files[idx].stem","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-17T07:59:13.044442Z","iopub.execute_input":"2025-04-17T07:59:13.044800Z","iopub.status.idle":"2025-04-17T07:59:13.050939Z","shell.execute_reply.started":"2025-04-17T07:59:13.044775Z","shell.execute_reply":"2025-04-17T07:59:13.049975Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load best model and predict\ntest_files = list(Path('/kaggle/input/waveform-inversion/test').glob('*.npy'))\ntest_ds = TestDataset(test_files)\ntest_loader = DataLoader(test_ds, batch_size = Config.batch_size)\n#model.load_state_dict(torch.load('/kaggle/input/cnn-model-2/best_model_2.pth'))\n#model.eval()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Generate submission file\nx_cols = [f'x_{i}' for i in range(1, 70, 2)]\nfieldnames = ['oid_ypos'] + x_cols\nwith open('submission.csv', 'w') as f:\n    writer = csv.DictWriter(f, fieldnames = fieldnames)\n    writer.writeheader()\n    with torch.no_grad():\n        for X, oids in tqdm(test_loader, desc = 'Predicting'):\n            X = X.to(device)\n            outputs = model(X).squeeze(1).cpu().numpy()\n            # Denormalize predictions\n            outputs = outputs * output_std + output_mean\n            # Format predictions to submission specs\n            for oid, pred in zip(oids, outputs):\n                for y_pos in range(Config.output_shape[0]):\n                    row = {'oid_ypos': f'{oid}_y_{y_pos}'}\n                    # Select odd-indexed columns\n                    row.update({f'x_{x}': pred[y_pos, x] for i, x in enumerate(range(1, Config.output_shape[1], 2))})\n                    writer.writerow(row)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}