add app

Dockerfile

@@ -0,0 +1,18 @@
+# Use an official Python runtime
+FROM python:3.9
+
+# Set the working directory
+WORKDIR /app
+
+# Copy all files to the container
+COPY . /app
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Expose the port Dash runs on
+EXPOSE 7860
+
+# Command to run the app
+CMD ["python", "app.py"]
+

app.py

@@ -0,0 +1,101 @@
+import os
+import numpy as np
+import plotly.graph_objects as go
+from dash import Dash, dcc, html, Input, Output
+import dash_bootstrap_components as dbc
+from predict import load_model, Configs
+import jax.numpy as jnp
+
+# Initialize Dash with a dark theme
+app = Dash(__name__, external_stylesheets=[dbc.themes.DARKLY])
+
+nn = Configs.nn
+x = np.linspace(-0.5, 0.5, nn, endpoint=True)
+y = np.linspace(-0.5, 0.5, nn, endpoint=True)
+
+forward1 = load_model('RVE1')
+forward2 = load_model('RVE2')
+forward3 = load_model('RVE3')
+
+# Nord color palette
+nord_bg = "#2E3440"  # Background
+nord_fg = "#D8DEE9"  # Text color
+nord_grid = "#4C566A"  # Grid color
+colormap = 'RdBu'  # Better for dark themes
+
+# PDE Solver
+def solve_pde(boundary_conditions, rve):
+    x = jnp.array(boundary_conditions)
+    return forward1(x) if rve == 'RVE1' else forward2(x) if rve == 'RVE2' else forward3(x)
+
+initial_bc = [0.02, 0.02, 0.02]
+d, s = solve_pde(initial_bc, 'RVE1')
+
+# Function to create Heatmap figures with a dark theme
+def create_figure(z_data, title):
+    fig = go.Figure(go.Contour(x=x, y=y, z=z_data, colorscale=colormap, showscale=False, line_smoothing=1.2, line_width=0))
+    fig.update_layout(
+        title=title, title_font=dict(color=nord_fg, size=18),
+        paper_bgcolor=nord_bg, plot_bgcolor=nord_bg, font=dict(color=nord_fg),
+        xaxis=dict(gridcolor=nord_grid, tickvals=[], scaleanchor="y", showgrid=False, zeroline=False, showline=False),
+        yaxis=dict(gridcolor=nord_grid, tickvals=[], showgrid=False, zeroline=False, showline=False),
+        autosize=False, width=300, height=300, margin=dict(l=0, r=0, t=30, b=0)
+    )
+    return fig
+
+titles = ["Disp. x", "Disp. y", "Stress xx", "Stress yy", "Stress xy"]
+pde_figs = [create_figure(d[:,:,0], titles[0]), create_figure(d[:,:,1], titles[1]),
+            create_figure(s[:,:,0], titles[2]), create_figure(s[:,:,1], titles[3]),
+            create_figure(s[:,:,2], titles[4])]
+
+# Layout
+app.layout = dbc.Container([
+    html.H1("RVE simulation using Equilibrium Neural Operator (EquiNO)", style={'textAlign': 'center', 'color': nord_fg, 'padding': '100px'}),
+    
+    dbc.Row([
+        dbc.Col([
+            html.H3("Prescribed Global Strains", style={'textAlign': 'center', 'color': nord_fg, 'padding': '10px'}),
+            html.Label("Strain xx", style={'font-size': '18px', 'color': nord_fg}),
+            dcc.Slider(id='bc1-slider', min=-0.04, max=0.04, step=0.005, value=initial_bc[0], marks=None,
+                       tooltip={"placement": "top", "always_visible": True}, updatemode='drag',
+                       className='slider-no-border'),
+            
+            html.Label("Strain yy", style={'font-size': '18px', 'color': nord_fg}),
+            dcc.Slider(id='bc2-slider', min=-0.04, max=0.04, step=0.005, value=initial_bc[1], marks=None,
+                       tooltip={"placement": "top", "always_visible": True}, updatemode='drag',
+                       className='slider-no-border'),
+            
+            html.Label("Strain xy", style={'font-size': '18px', 'color': nord_fg}),
+            dcc.Slider(id='bc3-slider', min=-0.04, max=0.04, step=0.005, value=initial_bc[2], marks=None,
+                       tooltip={"placement": "top", "always_visible": True}, updatemode='drag',
+                       className='slider-no-border'),
+            
+            html.Label("RVE Selector", style={'font-size': '18px', 'color': nord_fg}),
+            dcc.Dropdown(id='rve-selector', options=[
+                {'label': 'RVE1', 'value': 'RVE1'}, {'label': 'RVE2', 'value': 'RVE2'}, {'label': 'RVE3', 'value': 'RVE3'}
+            ], value='RVE1', clearable=False, style={'backgroundColor': 'white', 'color': nord_bg}),
+        ], width=3, style={'padding': '100px', 'backgroundColor': nord_bg}),
+        
+        dbc.Col([
+            dbc.Row([dcc.Graph(id=f'pde-figure{i+1}', figure=pde_figs[i], style={'width': '19%', 'display': 'inline-block'}) for i in range(5)]),
+        ], width=9, style={'padding': '100px', 'backgroundColor': nord_bg})
+    ])
+], fluid=True, style={'backgroundColor': nord_bg, 'minHeight': '200vh'})
+
+# Callback
+@app.callback(
+    [Output(f'pde-figure{i+1}', 'figure') for i in range(5)],
+    [Input('bc1-slider', 'value'), Input('bc2-slider', 'value'), Input('bc3-slider', 'value'), Input('rve-selector', 'value')]
+)
+def update_pde(bc1, bc2, bc3, selected_rve):
+    new_bc = [bc1, bc2, bc3]
+    d, s = solve_pde(new_bc, selected_rve)
+    return [create_figure(d[:,:,0], titles[0]), create_figure(d[:,:,1], titles[1]),
+            create_figure(s[:,:,0], titles[2]), create_figure(s[:,:,1], titles[3]),
+            create_figure(s[:,:,2], titles[4])]
+
+# Run the app
+server = app.server
+if __name__ == "__main__":
+    app.run_server(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))
+

predict.py

@@ -0,0 +1,97 @@
+import numpy as np
+import tensorflow as tf
+import jax
+import jax.numpy as jnp
+from scipy.interpolate import LinearNDInterpolator
+
+class Configs:
+    nn = 101
+
+def load_model(rve):
+    mdir = f'./trained_models/EquiNO_{rve}'
+    model = tf.keras.models.load_model(mdir)
+    nodes = np.array(model.kinema.get_config()['nodeCoord'])
+    scaling_bu = model.get_config()['scaling_bu']
+    scaling_bs = model.get_config()['scaling_bs']
+    scale_input = 0.04
+    
+    shp = np.array(model.kinema.get_config()['shp'])
+    elemInc = np.array(model.kinema.get_config()['elemInc'])
+    
+    # Extract model outputs and reshape arrays
+    t, _, t_s = model.trunk(nodes)
+    t = t.numpy()
+    t_s = t_s.numpy()
+    t_s = t_s.reshape((-1, 9, 16, 3))
+    t_s = jnp.einsum('nlrk,lj->njrk', t_s, jnp.linalg.pinv(shp[:, 0]))
+    
+    t_sg = np.zeros((nodes.shape[0], 16, 3))
+    t_sg[elemInc[:, 2:] - 1] = np.array(t_s)
+    t_s = t_sg
+    
+    x = np.linspace(-0.5, 0.5, Configs.nn, endpoint=True)
+    y = np.linspace(-0.5, 0.5, Configs.nn, endpoint=True)
+    xx, yy = np.meshgrid(x, y)
+    
+    def interp(d, s):
+        d = d.reshape((-1, 16*2))
+        s = s.reshape((-1, 16*3))
+        d = LinearNDInterpolator(nodes, d)(xx, yy)
+        s = LinearNDInterpolator(nodes, s)(xx, yy)
+        d = d.reshape((-1, 16, 2))
+        s = s.reshape((-1, 16, 3))
+        return d, s
+    
+    t, t_s = interp(t, t_s)
+    
+    nodes = np.stack((xx.flatten(), yy.flatten()), 1)
+    
+    weights = model.branch.get_weights()
+
+    weights_u = [weights[i:i+2] for i in range(0, len(weights), 4)]
+    params_u = [(jnp.array(w[0]), jnp.array(w[1])) for w in weights_u]
+    
+    weights_s = [weights[i:i+2] for i in range(2, len(weights), 4)]
+    params_s = [(jnp.array(w[0]), jnp.array(w[1])) for w in weights_s]
+    
+    cd = jnp.array([[2.0, 0.0], [0.0, 2.0]])
+
+    del model
+    
+    @jax.jit
+    def periodic_disp(x):
+        matrix = nodes[..., None] * cd[None, ...]
+        matrix = 0.5 * jnp.concatenate([matrix, jnp.flip(nodes, 1)[:, None, :]], 1)
+        return jnp.einsum('ij,ljm->ilm', x, matrix)
+
+    @jax.jit
+    def jax_branch_s(x):
+        x_n = x
+        for (w, b) in params_s[:-1]:
+            x_n = jax.nn.swish(jnp.dot(x_n, w) + b)
+        final_w, final_b = params_s[-1]
+        return jnp.dot(x_n, final_w) + final_b
+    
+    @jax.jit
+    def jax_branch_u(x):
+        x_n = x
+        for (w, b) in params_u[:-1]:
+            x_n = jax.nn.swish(jnp.dot(x_n, w) + b)
+        final_w, final_b = params_u[-1]
+        return jnp.dot(x_n, final_w) + final_b
+    
+    @jax.jit
+    def forward(x):
+        x = x.reshape((1, 3))
+        b = jax_branch_u(x / scale_input)
+        b_s = jax_branch_s(x / scale_input)
+        b = b * scaling_bu[1] + scaling_bu[0]
+        b_s = b_s * scaling_bs[1] + scaling_bs[0]
+        u = jnp.einsum('im,lmn->iln', b, t)
+        s = jnp.einsum('im,lmn->iln', b_s, t_s)
+        u = u - u[:, :1] + periodic_disp(x)
+        u = u.reshape(Configs.nn, Configs.nn, -1)
+        s = s.reshape(Configs.nn, Configs.nn, -1)
+        return u, s
+    
+    return forward

requirements.txt

@@ -0,0 +1,10 @@
+dash
+dash_bootstrap_components
+jax==0.4.20
+jaxlib==0.4.20
+numpy==1.24.3
+plotly==5.9.0
+scipy==1.10.1
+tensorflow==2.14.0
+Flask
+gunicorn