<commit_before>import numpy as np
import numpy.ma as ma
import rasterio
from ncdjango.geoprocessing.params import StringParameter
from ncdjango.geoprocessing.workflow import Task
from ncdjango.models import Service, Variable, SERVICE_DATA_ROOT
from django.conf import settings
import os.path
import tempfile
import netCDF4


class WriteTIF(Task):
    name = 'write_tif'
    inputs = [
        StringParameter('service_id'),
    ]
    outputs = [StringParameter('filename')]

    def execute(self, service_id):
        svc = Service.objects.get(name=service_id)
        var = Variable.objects.get(service_id=svc.id)
        data_path = svc.data_path
        with netCDF4.Dataset(os.path.join(SERVICE_DATA_ROOT, data_path), 'r') as nc:
            data = nc.variables[var.name][:]

        height, width = data.shape
        ex = var.full_extent
        x_step = (ex.xmax - ex.xmin) / width
        y_step = (ex.ymax - ex.ymin) / height
        transform = [ex.xmin, x_step, 0, ex.ymax, 0, -y_step]
        dtype = np.int16
        nodata = data.fill_value

        fd, filename = tempfile.mkstemp(prefix=settings.DATASET_DOWNLOAD_DIR, suffix='.tif')
        os.close(fd)
        with rasterio.Env(GDAL_TIFF_INTERNAL_MASK=True) as env:
            with rasterio.open(filename, 'w', driver='GTiff',
              height=height, width=width,
              crs=var.projection, transform=transform,
              count=1, dtype=dtype, nodata=nodata) as dst:
                dst.write(np.array(data, dtype=dtype), 1)
                dst.write_mask(np.logical_not(data.mask))

        return os.path.basename(filename)<commit_msg>Compress GeoTIFF export, encode as uint8<commit_after>import zipfile

import netCDF4
import numpy as np
import os.path
import rasterio
import tempfile
from django.conf import settings
from io import BytesIO
from ncdjango.geoprocessing.params import StringParameter
from ncdjango.geoprocessing.workflow import Task
from ncdjango.models import Service, Variable, SERVICE_DATA_ROOT
from rasterio.transform import Affine


class WriteTIF(Task):
    name = 'write_tif'
    inputs = [
        StringParameter('service_id')
    ]
    outputs = [StringParameter('filename')]

    def execute(self, service_id):
        svc = Service.objects.get(name=service_id)
        var = Variable.objects.get(service_id=svc.id)
        data_path = svc.data_path
        with netCDF4.Dataset(os.path.join(SERVICE_DATA_ROOT, data_path), 'r') as nc:
            data = nc.variables[var.name][:].astype('uint8')

        height, width = data.shape
        ex = var.full_extent
        x_step = (ex.xmax - ex.xmin) / width
        y_step = (ex.ymax - ex.ymin) / height
        transform = Affine.from_gdal(ex.xmin, x_step, 0, ex.ymax, 0, -y_step)
        dtype = np.uint8
        nodata = data.fill_value

        fd, filename = tempfile.mkstemp(prefix=settings.DATASET_DOWNLOAD_DIR, suffix='.zip')
        os.close(fd)

        with zipfile.ZipFile(filename, mode='w') as zf:
            tif_data = BytesIO()

            with rasterio.Env(GDAL_TIFF_INTERNAL_MASK=True):
                opts = dict(
                    driver='GTiff', height=height, width=width, crs=var.projection, transform=transform, count=1,
                    dtype=dtype, nodata=nodata
                )
                with rasterio.open(tif_data, 'w', **opts) as dst:
                    dst.write(np.array(data, dtype=dtype), 1)
                    dst.write_mask(np.logical_not(data.mask))

                zf.writestr('sst_results.tif', tif_data.getvalue(), compress_type=zipfile.ZIP_DEFLATED)

        return os.path.basename(filename)
{'fraction_non_alphanumeric': 0.10542335362479248, 'fraction_numerical': 0.004703929164360819, 'mean_word_length': 2.685015290519878, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2998851', 'n_tokens_mistral': 1207, 'n_tokens_neox': 1170, 'n_words': 243}