How to add new inputs and features in model FNO 2d

Abdulla_Kh: <p>Hi everyone!<br>
Guys, I am new in pytorch and i try to adopt model FNO2d(ref: <a href="https://github.com/neuraloperator/neuraloperator/blob/master/fourier_2d_time.py" class="inline-onebox" rel="noopener nofollow ugc">neuraloperator/fourier_2d_time.py at master · neuraloperator/neuraloperator · GitHub</a>)</p>
<p>Task:<br>
I’m trying to predict the pressure field in oil production.<br>
I have two wells for extraction and injection. In its primary form, the model FNO includes only a field without properties(additional inputs) and I’m trying to add these properties like: porosity, permeability and well location.</p>
<p>First training data set had size [Case = 300,dim1 = 64, dim2 = 86, timesteps= 84], after adding new properties into data size changed to [300, 64, 86, 84, 4] (under 0 -  pressure feld).</p>
<p>And this is how I did it:</p>
<pre><code class="lang-auto">        self.modes1 = modes1
        self.modes2 = modes2
        self.width = width
        self.padding = 8 # pad the domain if input is non-periodic

        self.p = nn.Linear(15, self.width) # input channel is 15: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y) + 3 properties
        self.conv0 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.conv1 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.conv2 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.conv3 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.mlp0 = MLP(self.width, self.width, self.width)
        self.mlp1 = MLP(self.width, self.width, self.width)
        self.mlp2 = MLP(self.width, self.width, self.width)
        self.mlp3 = MLP(self.width, self.width, self.width)
        self.w0 = nn.Conv2d(self.width, self.width, 1)
        self.w1 = nn.Conv2d(self.width, self.width, 1)
        self.w2 = nn.Conv2d(self.width, self.width, 1)
        self.w3 = nn.Conv2d(self.width, self.width, 1)
        self.norm = nn.InstanceNorm2d(self.width)
        self.q = MLP(self.width, 1, self.width * 4) # output channel is 1: u(x, y)

    def forward(self, x):

        grid = self.get_grid(x.shape, x.device)
        pressure = self.get_Pressure(x, x.shape, x.device)
        properties = self.get_Properties(x, x.shape, x.device)

        x = torch.cat((pressure, grid, properties), dim=-1)
        x = self.p(x)
        x = x.permute(0, 3, 1, 2)
#....and so on
</code></pre>
<p>But it doesn’t work correctly. When i try predict pressure with new well location, it has too big loss and visual also not correct. But when wells location same like in training data(in the training dataset wells had fixed location for all cases) it work correctly with some noises.<br>
Example:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/2/226f1d9b3ffd8aa55bc601b8b07cabd8dd41d44b.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/226f1d9b3ffd8aa55bc601b8b07cabd8dd41d44b" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/226f1d9b3ffd8aa55bc601b8b07cabd8dd41d44b_2_690x376.jpeg" alt="image" data-base62-sha1="4UCiuc7Omkh8FNOxskrfUYkactB" width="690" height="376" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/226f1d9b3ffd8aa55bc601b8b07cabd8dd41d44b_2_690x376.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/226f1d9b3ffd8aa55bc601b8b07cabd8dd41d44b_2_1035x564.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/2/226f1d9b3ffd8aa55bc601b8b07cabd8dd41d44b.jpeg 2x" data-dominant-color="B7DFB8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1273×694 48.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I would be grateful for any help or advice.</p>