Update only Z values in a mapbox?

Hi, I’m new to Dash and find it absolutely amazing.

My question: I have a Choroplethmapbox and a series of 10 buttons that switch between map modes. All modes are based on the same geojson and mapbox, and the buttons change only the z values (and the hover data). The z’s are just 10 different columns in the same dataframe: button A switches the z-values to column A, B - to column B, etc. All works perfectly, but as expected, the switch is quite slow - probably because I have to send the whole figure with every callback (a similar question: Updating data on mapbox without updating the whole map).

Ideally I would like to update just the specific property (as suggested here: Is it possible to update just `layout`, not whole `figure` of Graph in callback?): in my case something like @app.callback(Output('my-map', 'figure['data'][0]['z']')). But as far as I understand it’s impossible. Is there any other solution for that? Can caching help here in any way? I’m attaching a simplified version of my code below. I’d appreciate your comments!

# Data

map_modes = { ... }
button_ids = list(map_modes.keys())

with open( ... ) as f: 
   my_geojson = json.load(f)

my_df = pd.read_csv( ... )
my_df.set_index('id', inplace=True)

data = [
	go.Choroplethmapbox(
		geojson=my_geojson,
		locations=my_df.index,
		z=my_df['A'],
        )
]

layout = go.Layout(
	mapbox={
		'style': 'carto-positron',
                ...
	},
	uirevision='same same'
)

# Layout

app.layout = html.Div(className='main_disp', children=[
	dcc.Graph(
		id='main_map'
	),
	html.Div(className='mode_buttons', 
		children=[html.Button(map_modes[x]['label'], id=x, className='inactive_button') for x in map_modes]
	),
])

# Callbacks

@app.callback(
	[Output('main_map', 'figure')] + [Output(button, 'className') for button in map_modes],
	[Input(button, 'n_clicks') for button in map_modes]
)
def button_pressed(*args):
    # when initialized, first button turns active, map remains default
	if not any(args):
		buttons_output = ['active_button'] + (['inactive_button'] * (len(button_ids) - 1))
	# otherwise, the fired button turns active and the button's column in the df becomes z
    else:
		ctx = dash.callback_context
		fired_button = ctx.triggered[0]['prop_id'].split('.')[0]
		data[0]['z'] = my_df[fired_button]

	figure_output = {'data': data, 'layout': layout}
	buttons_output = ['inactive_button' if button != fired_button else 'active_button' for button in button_ids]
	return [figure_output] + buttons_output