<commit_before># -*- coding: utf-8 -*-
"""
==============================================================================
View mixin classes (:mod:`sknano.apps.nanogen_gui._view_mixins`)
==============================================================================

.. currentmodule:: sknano.apps.nanogen_gui._view_mixins

"""
from __future__ import absolute_import, division, print_function, \
    unicode_literals
__docformat__ = 'restructuredtext en'

try:
    # from PyQt4.QtCore import pyqtSlot
    from PyQt5.QtCore import pyqtSlot
except ImportError as e:
    print(e)

__all__ = ['NanotubeViewMixin', 'GrapheneViewMixin']

class NanotubeViewMixin:
    """Mixin class for nanotube classes."""
    @pyqtSlot()
    def on_n_spinBox_editingFinished(self):
        self.model.n = self.n_spinBox.value()

    @pyqtSlot(int)
    def on_n_spinBox_valueChanged(self, value):
        self.model.n = value

    @pyqtSlot()
    def on_m_spinBox_editingFinished(self):
        self.model.m = self.m_spinBox.value()

    @pyqtSlot(int)
    def on_m_spinBox_valueChanged(self, value):
        self.model.m = value

    @pyqtSlot()
    def on_nanotube_bond_doubleSpinBox_editingFinished(self):
        self.model.nanotube_bond = self.nanotube_bond_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_nanotube_bond_doubleSpinBox_valueChanged(self, value):
        self.model.nanotube_bond = value

    @pyqtSlot()
    def on_nx_spinBox_editingFinished(self):
        self.model.nx = self.nx_spinBox.value()

    @pyqtSlot(int)
    def on_nx_spinBox_valueChanged(self, value):
        self.model.nx = value

    @pyqtSlot()
    def on_ny_spinBox_editingFinished(self):
        self.model.ny = self.ny_spinBox.value()

    @pyqtSlot(int)
    def on_ny_spinBox_valueChanged(self, value):
        self.model.ny = value

    @pyqtSlot()
    def on_nz_spinBox_editingFinished(self):
        self.model.nz = self.nz_spinBox.value()

    @pyqtSlot(int)
    def on_nz_spinBox_valueChanged(self, value):
        self.model.nz = value

    @pyqtSlot(int)
    def on_nanotube_generator_buttonGroup_buttonClicked(self):
        if self.unrolled_nanotube_generator_radioButton.isChecked():
            if self.bundle_generator_checkBox.isChecked():
                self.bundle_generator_checkBox.setCheckState(False)
            self.bundle_generator_checkBox.setCheckable(False)
        else:
            self.bundle_generator_checkBox.setCheckable(True)

    @pyqtSlot(str)
    def on_nanotube_element1_comboBox_currentIndexChanged(self, value):
        self.model.nanotube_element1 = str(value)

    @pyqtSlot(str)
    def on_nanotube_element2_comboBox_currentIndexChanged(self, value):
        self.model.nanotube_element2 = str(value)

    @pyqtSlot(str)
    def on_graphene_element1_comboBox_currentIndexChanged(self, value):
        self.model.graphene_element1 = str(value)

    @pyqtSlot(str)
    def on_graphene_element2_comboBox_currentIndexChanged(self, value):
        self.model.graphene_element2 = str(value)

    @pyqtSlot()
    def on_Lz_doubleSpinBox_editingFinished(self):
        self.model.Lz = self.Lz_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_Lz_doubleSpinBox_valueChanged(self, value):
        self.model.Lz = value

class GrapheneViewMixin(object):
    """Mixin class for graphene."""

    # @pyqtSlot(int)
    # def on_stacking_order_buttonGroup_buttonClicked(self):
    #     if self.nlayer_spinBox.value() == 1:
    #         pass

    @pyqtSlot()
    def on_armchair_edge_length_doubleSpinBox_editingFinished(self):
        self.model.armchair_edge_length = \
            self.armchair_edge_length_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_armchair_edge_length_doubleSpinBox_valueChanged(self, value):
        self.model.armchair_edge_length = value

    @pyqtSlot()
    def on_zigzag_edge_length_doubleSpinBox_editingFinished(self):
        self.model.zigzag_edge_length = \
            self.zigzag_edge_length_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_zigzag_edge_length_doubleSpinBox_valueChanged(self, value):
        self.model.zigzag_edge_length = value

    @pyqtSlot()
    def on_nlayers_spinBox_editingFinished(self):
        self.model.nlayers = self.nlayers_spinBox.value()
        #self.update_stacking_order_buttonGroup()

    @pyqtSlot(int)
    def on_nlayers_spinBox_valueChanged(self, value):
        self.model.nlayers = value
        # self.update_stacking_order_buttonGroup()

    # def update_stacking_order_buttonGroup(self):
    #     if self.nlayers_spinBox.value() == 1:
    #         for rb in (self.AA_stacking_radioButton,
    #                    self.AB_stacking_radioButton):
    #             rb.setCheckable(False)
    #     else:
    #         for rb in (self.AA_stacking_radioButton,
    #                    self.AB_stacking_radioButton):
    #             rb.setCheckable(True)

    @pyqtSlot()
    def on_graphene_bond_doubleSpinBox_editingFinished(self):
        self.model.graphene_bond = self.graphene_bond_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_graphene_bond_doubleSpinBox_valueChanged(self, value):
        self.model.graphene_bond = value
<commit_msg>adding new mixin classes<commit_after># -*- coding: utf-8 -*-
"""
==============================================================================
View mixin classes (:mod:`sknano.apps.nanogen_gui._view_mixins`)
==============================================================================

.. currentmodule:: sknano.apps.nanogen_gui._view_mixins

"""
from __future__ import absolute_import, division, print_function, \
    unicode_literals
__docformat__ = 'restructuredtext en'

try:
    # from PyQt4.QtCore import pyqtSlot
    from PyQt5.QtCore import pyqtSlot
except ImportError as e:
    print(e)

__all__ = ['MainWindowViewMixin', 'SWNTViewMixin', 'MWNTViewMixin',
           'GrapheneViewMixin', 'FullereneViewMixin']

class MainWindowViewMixin:
    """Mixin class for main window."""

    @pyqtSlot(str)
    def on_element1_comboBox_currentIndexChanged(self, value):
        self.model.element1 = str(value)

    @pyqtSlot(str)
    def on_element2_comboBox_currentIndexChanged(self, value):
        self.model.element2 = str(value)

    @pyqtSlot()
    def on_bond_doubleSpinBox_editingFinished(self):
        self.model.bond = self.bond_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_bond_doubleSpinBox_valueChanged(self, value):
        self.model.bond = value

class SWNTViewMixin:
    """Mixin class for nanotube classes."""
    @pyqtSlot()
    def on_swnt_n_spinBox_editingFinished(self):
        self.model.n = self.swnt_n_spinBox.value()

    @pyqtSlot(int)
    def on_swnt_n_spinBox_valueChanged(self, value):
        self.model.n = value

    @pyqtSlot()
    def on_swnt_m_spinBox_editingFinished(self):
        self.model.m = self.swnt_m_spinBox.value()

    @pyqtSlot(int)
    def on_swnt_m_spinBox_valueChanged(self, value):
        self.model.m = value

    @pyqtSlot()
    def on_swnt_nz_doubleSpinBox_editingFinished(self):
        self.model.nz = self.swnt_nz_doubleSpinBox.value()

    @pyqtSlot(int)
    def on_swnt_nz_doubleSpinBox_valueChanged(self, value):
        self.model.nz = value

    @pyqtSlot()
    def on_swnt_Lz_doubleSpinBox_editingFinished(self):
        self.model.Lz = self.swnt_Lz_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_swnt_Lz_doubleSpinBox_valueChanged(self, value):
        self.model.Lz = value

    @pyqtSlot()
    def on_swnt_bundle_nx_spinBox_editingFinished(self):
        self.model.nx = self.nx_spinBox.value()

    @pyqtSlot(int)
    def on_swnt_bundle_nx_spinBox_valueChanged(self, value):
        self.model.nx = value

    @pyqtSlot()
    def on_swnt_bundle_ny_spinBox_editingFinished(self):
        self.model.ny = self.ny_spinBox.value()

    @pyqtSlot(int)
    def on_swnt_bundle_ny_spinBox_valueChanged(self, value):
        self.model.ny = value

    # @pyqtSlot(int)
    # def on_swnt_tab_generator_buttonGroup_buttonClicked(self):
    #     if self.unrolled_swnt_generator_radioButton.isChecked():
    #         if self.bundle_generator_checkBox.isChecked():
    #             self.bundle_generator_checkBox.setCheckState(False)
    #         self.bundle_generator_checkBox.setCheckable(False)
    #     else:
    #         self.bundle_generator_checkBox.setCheckable(True)

class MWNTViewMixin:

    @pyqtSlot()
    def on_mwnt_bundle_nx_spinBox_editingFinished(self):
        self.model.nx = self.mwnt_bundle_nx_spinBox.value()

    @pyqtSlot(int)
    def on_mwnt_bundle_nx_spinBox_valueChanged(self, value):
        self.model.nx = value

    @pyqtSlot()
    def on_mwnt_bundle_ny_spinBox_editingFinished(self):
        self.model.ny = self.mwnt_bundle_ny_spinBox.value()

    @pyqtSlot(int)
    def on_mwnt_bundle_ny_spinBox_valueChanged(self, value):
        self.model.ny = value

    @pyqtSlot()
    def on_mwnt_Lz_spinBox_editingFinished(self):
        self.model.Lz = self.mwnt_Lz_doubleSpinBox.value()

    @pyqtSlot(int)
    def on_mwnt_Lz_doubleSpinBox_valueChanged(self, value):
        self.model.Lz = value

class GrapheneViewMixin:
    """Mixin class for graphene."""

    # @pyqtSlot(int)
    # def on_stacking_order_buttonGroup_buttonClicked(self):
    #     if self.nlayer_spinBox.value() == 1:
    #         pass

    @pyqtSlot()
    def on_armchair_edge_length_doubleSpinBox_editingFinished(self):
        self.model.armchair_edge_length = \
            self.armchair_edge_length_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_armchair_edge_length_doubleSpinBox_valueChanged(self, value):
        self.model.armchair_edge_length = value

    @pyqtSlot()
    def on_zigzag_edge_length_doubleSpinBox_editingFinished(self):
        self.model.zigzag_edge_length = \
            self.zigzag_edge_length_doubleSpinBox.value()

    @pyqtSlot(float)
    def on_zigzag_edge_length_doubleSpinBox_valueChanged(self, value):
        self.model.zigzag_edge_length = value

    @pyqtSlot()
    def on_nlayers_spinBox_editingFinished(self):
        self.model.nlayers = self.nlayers_spinBox.value()
        #self.update_stacking_order_buttonGroup()

    @pyqtSlot(int)
    def on_nlayers_spinBox_valueChanged(self, value):
        self.model.nlayers = value
        # self.update_stacking_order_buttonGroup()

    # def update_stacking_order_buttonGroup(self):
    #     if self.nlayers_spinBox.value() == 1:
    #         for rb in (self.AA_stacking_radioButton,
    #                    self.AB_stacking_radioButton):
    #             rb.setCheckable(False)
    #     else:
    #         for rb in (self.AA_stacking_radioButton,
    #                    self.AB_stacking_radioButton):
    #             rb.setCheckable(True)

class FullereneViewMixin:
    pass
<|endoftext|>


<commit_before>import os
import sys
import vcf
import gzip
import pandas as pd
from collections import defaultdict

from argparse import ArgumentParser

def _moderate(record):
    if record.find("MODERATE") > -1 or record.find("HIGH") > -1:
        return True

def _gtp(info):
    return info.split(":")[0]

def _info_sample(name, info, tag):
    k = name.split(":")
    v = info.split(":")
    d = dict(zip(k, v))
    return d[tag]

def _is_in_any_control(samples, t, n):
    c = set()
    for s in n:
        if s == "./.":
            continue
        c.add(s)
        c.add(s[::-1])
    skip = 0
    for s in t:
        if _gtp(samples[s]) in c:
           skip += 1
    return skip

def _get_info(info):
    cols = info.split(";")
    effect = []
    genes = []
    for field in cols:
        if field.find("=") == -1:
            continue
        k, v = field.split("=")
        if k == "STATUS":
            status = v
        elif k == "EFF":
            for e in v.split(","):
                type_e = e.split("(")[0]
                gene_e = e.split("(")[1].split("|")[5]
                level_e = e.split("(")[1].split("|")[0]
                aa_e = e.split("(")[1].split("|")[3]
                effect.append("|".join([type_e, level_e, aa_e]))
                genes.append(gene_e)
    return [status, effect, genes]

def _frmt(info, t, n, genes=None):
    """Get information for the variant"""
    in_gene = "No"
    cols = info.split("\t")
    samples = cols[9:]
    chrom, pos, rs = cols[:3]
    change = "%s:%s" % (cols[3], cols[4])
    status, eff, gene = _get_info(cols[7])
    effect = ";".join(list(set(["%s=%s" % (e, g) for e, g in zip(eff, gene)])))
    if any([g in genes for g in gene]):
        in_gene = "Yes"
    gt_t = ":".join([_gtp(samples[s]) for s in t])
    gt_n = ":".join([_gtp(samples[s]) for s in n])
    af_t = ":".join([_info_sample(cols[8], samples[s], "AF") for s in t])
    af_n = ":".join([_info_sample(cols[8], samples[s], "AF") for s in n])
    line = ("{chrom}\t{pos}\t{rs}\t{change}\t{gt_t}|{gt_n}\t{af_t}|{af_n}\t{status}\t{effect}\t{in_gene}")
    return line.format(**locals())

def _read_genes(fn):
    if fn:
        df = pd.read_csv(fn)
        return set(df["Gene"])

def _add_to_genes(record, t, n, d):
    cols = record.strip().split("\t")
    samples = cols[9:]
    status, eff, gene = _get_info(cols[7])
    for s in [0, 1, 2]:
        if not _gtp(samples[t[s]]) == "./." and not _gtp(samples[n[s]]) == "./.":
            if _gtp(samples[t[s]]) != _gtp(samples[n[s]]):
                for g in gene:
                    if g not in d:
                        d[g].update({s: []})
                    if s not in d[g]:
                        d[g].update({s: []})
                    d[g][s].append(record)
    return d

def _read_vcf(fn, n):
    genotypes = defaultdict(list)
    with gzip.open(fn, 'rb') as vcf_reader:
        for record in vcf_reader:
            if not record.startswith("#"):
                samples = record.strip().split("\t")[9:]
                rs = "".join(record.strip().split("\t")[:2])
                for s in n:
                    genotypes[rs].append(_gtp(samples[s]))
    return genotypes

def _clean_variants(fn, genes, control):
    """
    Read variants and keep the ones in tumor
    """
    t = [0, 2, 4]
    n = [1, 3, 5]
    genes = _read_genes(genes)
    controls = _read_vcf(control, n)
    print "\t".join(["chrom", "pos", "rs", "change", "gt:N2N3N4|C2W3W4",
                     "af:tumor|normal", "status", "effects",
                     "in_gene_list", "num_samples_seen", "filter"])
    var_by_genes = defaultdict(dict)
    with gzip.open(fn[0], 'rb') as vcf_reader:
        for record in vcf_reader:
            skip = 0
            if not record.startswith("#"):
                samples = record.strip().split("\t")[9:]
                rs = "".join(record.strip().split("\t")[:2])
                for s in [0, 1, 2]:
                    skip += _gtp(samples[t[s]]) == "./." or _gtp(samples[n[s]]) == "./."
                reject = _is_in_any_control(samples, t, controls[rs])
                if reject > 0:
                    continue
                if skip > 1:
                    if _moderate(record):
                        var_by_genes = _add_to_genes(record, t, n, var_by_genes)
                frmt = _frmt(record, t, n, genes)
                #print "skip %s reject %s" % (skip, reject)
                print "%s\t%s\t%s" % (frmt, 3 - skip, "variant")
    var_seen = set()
    for g in var_by_genes:
        if len(var_by_genes[g]) > 2:
            for s in var_by_genes[g]:
                for var in var_by_genes[g][s]:
                    var_frmt = _frmt(var, t, n, genes)
                    if var_frmt not in var_seen:
                        print "%s\t%s\t%s" % (var_frmt, 1, "gene")
                        var_seen.add(var_frmt)

if __name__ == "__main__":
    parser = ArgumentParser(description="Detect enhancer activity.")
    parser.add_argument("--genes", help="List of genes to annotate.")
    parser.add_argument("--control", help="vcf with all genotypes.")
    parser.add_argument("files", nargs="*", help="vcf files.")
    args = parser.parse_args()

    clean_vcf = _clean_variants(args.files, args.genes, args.control)
<commit_msg>Update priorization<commit_after>import os
import sys
import vcf
import gzip
import pandas as pd
from collections import defaultdict

from argparse import ArgumentParser

def _moderate(record):
    if record.find("MODERATE") > -1 or record.find("HIGH") > -1:
        return True

def _gtp(info):
    return info.split(":")[0]

def _info_sample(name, info, tag):
    k = name.split(":")
    v = info.split(":")
    d = dict(zip(k, v))
    return d[tag]

def _is_in_any_control(samples, t, n):
    c = set()
    for s in n:
        if s == "1/1" or s.find("1") > -1:
            return 1
        if s == "./.":
            continue
        c.add(s)
        c.add(s[::-1])
    skip = 0
    for s in t:
        if _gtp(samples[s]) in c:
           skip += 1
    return skip

def _get_info(info):
    cols = info.split(";")
    effect = []
    genes = []
    for field in cols:
        if field.find("=") == -1:
            continue
        k, v = field.split("=")
        if k == "STATUS":
            status = v
        elif k == "EFF":
            for e in v.split(","):
                type_e = e.split("(")[0]
                gene_e = e.split("(")[1].split("|")[5]
                level_e = e.split("(")[1].split("|")[0]
                aa_e = e.split("(")[1].split("|")[3]
                effect.append("|".join([type_e, level_e, aa_e]))
                genes.append(gene_e)
    return [status, effect, genes]

def _frmt(info, t, n, genes=None):
    """Get information for the variant"""
    in_gene = "No"
    cols = info.split("\t")
    samples = cols[9:]
    chrom, pos, rs = cols[:3]
    change = "%s:%s" % (cols[3], cols[4])
    status, eff, gene = _get_info(cols[7])
    effect = ";".join(list(set(["%s=%s" % (e, g) for e, g in zip(eff, gene)])))
    if any([g in genes for g in gene]):
        in_gene = "Yes"
    gt_t = ":".join([_gtp(samples[s]) for s in t])
    gt_n = ":".join([_gtp(samples[s]) for s in n])
    af_t = ":".join([_info_sample(cols[8], samples[s], "AF") for s in t])
    af_n = ":".join([_info_sample(cols[8], samples[s], "AF") for s in n])
    line = ("{chrom}\t{pos}\t{rs}\t{change}\t{gt_t}|{gt_n}\t{af_t}|{af_n}\t{status}\t{effect}\t{in_gene}")
    return line.format(**locals())

def _read_genes(fn):
    if fn:
        df = pd.read_csv(fn)
        return set(df["Gene"])

def _add_to_genes(record, t, n, d):
    cols = record.strip().split("\t")
    samples = cols[9:]
    status, eff, gene = _get_info(cols[7])
    for s in [0, 1, 2]:
        if not _gtp(samples[t[s]]) == "./." and not _gtp(samples[n[s]]) == "./.":
            if _gtp(samples[t[s]]) != _gtp(samples[n[s]]):
                for g in gene:
                    if g not in d:
                        d[g].update({s: []})
                    if s not in d[g]:
                        d[g].update({s: []})
                    d[g][s].append(record)
    return d

def _read_vcf(fn, n):
    genotypes = defaultdict(list)
    with gzip.open(fn, 'rb') as vcf_reader:
        for record in vcf_reader:
            if not record.startswith("#"):
                samples = record.strip().split("\t")[9:]
                rs = "".join(record.strip().split("\t")[:2])
                for s in n:
                    genotypes[rs].append(_gtp(samples[s]))
    return genotypes

def _clean_variants(fn, genes, control):
    """
    Read variants and keep the ones in tumor
    """
    t = [0, 2, 4]
    n = [1, 3, 5]
    genes = _read_genes(genes)
    controls = _read_vcf(control, n)
    print "\t".join(["chrom", "pos", "rs", "change", "gt:N2N3N4|C2W3W4",
                     "af:tumor|normal", "status", "effects",
                     "in_gene_list", "num_samples_seen", "filter"])
    var_by_genes = defaultdict(dict)
    with gzip.open(fn[0], 'rb') as vcf_reader:
        for record in vcf_reader:
            skip = 0
            if not record.startswith("#"):
                samples = record.strip().split("\t")[9:]
                rs = "".join(record.strip().split("\t")[:2])
                for s in [0, 1, 2]:
                    skip += _gtp(samples[t[s]]) == "./." or _gtp(samples[n[s]]) == "./."
                reject = _is_in_any_control(samples, t, controls[rs])
                if reject > 0:
                    continue
                if skip > 1:
                    if _moderate(record):
                        var_by_genes = _add_to_genes(record, t, n, var_by_genes)

                frmt = _frmt(record, t, n, genes)
                # print "skip %s reject %s" % (skip, reject)
                print "%s\t%s\t%s" % (frmt, 3 - skip, "variant")

    var_seen = set()
    for g in var_by_genes:
        if len(var_by_genes[g]) >= 2:
            for s in var_by_genes[g]:
                for var in var_by_genes[g][s]:
                    var_frmt = _frmt(var, t, n, genes)
                    if var_frmt not in var_seen:
                        print "%s\t%s\t%s" % (var_frmt, len(var_by_genes[g]), "gene")
                        var_seen.add(var_frmt)

if __name__ == "__main__":
    parser = ArgumentParser(description="Detect enhancer activity.")
    parser.add_argument("--genes", help="List of genes to annotate.")
    parser.add_argument("--control", help="vcf with all genotypes.")
    parser.add_argument("files", nargs="*", help="vcf files.")
    args = parser.parse_args()

    clean_vcf = _clean_variants(args.files, args.genes, args.control)
<|endoftext|>


<commit_before>import sys
from django.core.management.base import BaseCommand
from ... import models

class Command(BaseCommand):

    help = "View and modify a document collection"

    def add_arguments(self, parser):
        parser.add_argument(
            'collection_slug',
            nargs='?',
            default=None,
            help="The slug of the collection to view/modify"
        )
        parser.add_argument(
            '--set-path',
            dest='path',
            default=None,
            help='Modify the path for the collection'
        )
        parser.add_argument(
            '--set-title',
            dest='title',
            default=None,
            help='Modify the user-readable title for the collection'
        )
        parser.add_argument(
            '--set-es-index',
            dest='es_index',
            default=None,
            help='Modify the Elasticsearch index for the collection'
        )
        parser.add_argument(
            '--set-slug',
            dest='slug',
            default=None,
            help='Modify the slug for the collection'
        )
        parser.add_argument(
            '--set-description',
            dest='description',
            default=None,
            help='Modify the single-line user-readable for the collection'
        )
        parser.add_argument(
            '--set-ocr',
            nargs=2,
            dest='ocr',
            default=None,
            help='[--set-ocr TAG PATH] Add or modify the path for an OCR file set'
        )
        parser.add_argument(
            '--remove-ocr',
            dest='remove_ocr',
            default=None,
            help='Remove a path for an OCR file set'
        )

    def handle(self, collection_slug, **options):
        if not collection_slug:
            print_all_collections()
            return

        try:
            collection = models.Collection.objects.filter(slug=collection_slug).get()
        except models.Collection.DoesNotExist:
            print("Collection with slug", collection_slug, "does not exist.")
            sys.exit(1)

        modify_collection(collection, **options)
        modify_ocr_data(collection, **options)
        print_info_for_collections([collection])

def print_all_collections():
    print_info_for_collections(models.Collection.objects.all().order_by('id'))

def print_info_for_collections(collections):
    for c in collections:
        print("===", c.id, c.slug, "===")
        print("title:", c.title)
        print("es_index:", c.es_index)
        print("description:", c.description)
        print("ocr:", c.ocr)
        print()

def modify_ocr_data(collection, **options):
    modified = False
    if options['ocr'] is not None:
        tag, path = options['ocr']
        collection.ocr[tag] = str(Path(path).resolve())
        modified = True
    if options['remove_ocr'] is not None:
        del collection.ocr[options['remove_ocr']]
        modified = True
    if modified:
        collection.save()

def modify_collection(collection, **options):
    keys = {
        'path',
        'slug',
        'es_index',
        'title',
        'description'
    }
    modified = False
    for key in keys:
        if options[key] is not None:
            setattr(collection, key, options[key])
            modified = True
    if modified:
        collection.save()

<commit_msg>refactor print loop<commit_after>import sys
from django.core.management.base import BaseCommand
from ... import models

class Command(BaseCommand):

    help = "View and modify a document collection"

    def add_arguments(self, parser):
        parser.add_argument(
            'collection_slug',
            nargs='?',
            default=None,
            help="The slug of the collection to view/modify"
        )
        parser.add_argument(
            '--set-path',
            dest='path',
            default=None,
            help='Modify the path for the collection'
        )
        parser.add_argument(
            '--set-title',
            dest='title',
            default=None,
            help='Modify the user-readable title for the collection'
        )
        parser.add_argument(
            '--set-es-index',
            dest='es_index',
            default=None,
            help='Modify the Elasticsearch index for the collection'
        )
        parser.add_argument(
            '--set-slug',
            dest='slug',
            default=None,
            help='Modify the slug for the collection'
        )
        parser.add_argument(
            '--set-description',
            dest='description',
            default=None,
            help='Modify the single-line user-readable for the collection'
        )
        parser.add_argument(
            '--set-ocr',
            nargs=2,
            dest='ocr',
            default=None,
            help='[--set-ocr TAG PATH] Add or modify the path for an OCR file set'
        )
        parser.add_argument(
            '--remove-ocr',
            dest='remove_ocr',
            default=None,
            help='Remove a path for an OCR file set'
        )

    def handle(self, collection_slug, **options):
        if not collection_slug:
            print_all_collections()
            return

        try:
            collection = models.Collection.objects.filter(slug=collection_slug).get()
        except models.Collection.DoesNotExist:
            print("Collection with slug", collection_slug, "does not exist.")
            sys.exit(1)

        modify_collection(collection, **options)
        modify_ocr_data(collection, **options)
        print_collection(collection)

def print_all_collections():
    collections = models.Collection.objects.all().order_by('id')
    for c in collections:
        print_collection(c)

def print_collection(c):
    print("===", c.id, c.slug, "===")
    print("title:", c.title)
    print("es_index:", c.es_index)
    print("description:", c.description)
    print("ocr:", c.ocr)
    print()

def modify_ocr_data(collection, **options):
    modified = False
    if options['ocr'] is not None:
        tag, path = options['ocr']
        collection.ocr[tag] = str(Path(path).resolve())
        modified = True
    if options['remove_ocr'] is not None:
        del collection.ocr[options['remove_ocr']]
        modified = True
    if modified:
        collection.save()

def modify_collection(collection, **options):
    keys = {
        'path',
        'slug',
        'es_index',
        'title',
        'description'
    }
    modified = False
    for key in keys:
        if options[key] is not None:
            setattr(collection, key, options[key])
            modified = True
    if modified:
        collection.save()

<|endoftext|>


<commit_before>#
# Copyright (C) 2007-2013 by frePPLe bvba
#
# This library is free software; you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero
# General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from datetime import datetime, timedelta

from django.contrib.admin.utils import unquote
from django.db import connections
from django.http import Http404
from django.utils.translation import ugettext_lazy as _

from freppledb.input.models import Demand
from freppledb.common.report import GridReport, GridFieldText, GridFieldNumber
from freppledb.common.models import Parameter

class ReportByDemand(GridReport):
  '''
  A list report to show peggings.
  '''
  template = 'output/pegging.html'
  title = _("Demand plan")
  filterable = False
  frozenColumns = 0
  editable = False
  default_sort = None
  hasTimeBuckets = True
  multiselect = False
  heightmargin = 87
  help_url = 'user-guide/user-interface/plan-analysis/demand-gantt-report.html'
  rows = (
    GridFieldText('depth', title=_('depth'), editable=False, sortable=False),
    GridFieldText('operation', title=_('operation'), editable=False, sortable=False, key=True, formatter='detail', extra='"role":"input/operation"'),
    GridFieldText('type', title=_('type'), editable=False, sortable=False, width=100),
    #GridFieldText('buffer', title=_('buffer'), formatter='buffer', editable=False, sortable=False),
    #GridFieldText('item', title=_('item'), formatter='item', editable=False, sortable=False),
    GridFieldText('resource', title=_('resource'), editable=False, sortable=False, extra='formatter:reslistfmt'),
    GridFieldNumber('quantity', title=_('quantity'), editable=False, sortable=False),
    GridFieldText('operationplans', width=1000, extra='formatter:ganttcell', editable=False, sortable=False),
    GridFieldText('parent', editable=False, sortable=False, hidden=True),
    GridFieldText('leaf', editable=False, sortable=False, hidden=True),
    GridFieldText('expanded', editable=False, sortable=False, hidden=True),
    GridFieldText('current', editable=False, sortable=False, hidden=True),
    GridFieldText('due', editable=False, sortable=False, hidden=True),
    GridFieldText('showdrilldown', editable=False, sortable=False, hidden=True),
    )

  @ classmethod
  def basequeryset(reportclass, request, args, kwargs):
    return Demand.objects.filter(name__exact=unquote(args[0])).values('name')

  @classmethod
  def getBuckets(reportclass, request, *args, **kwargs):
    # Get the earliest and latest operationplan, and the demand due date
    cursor = connections[request.database].cursor()
    cursor.execute('''
      with dmd as (
        select
          due,
          cast(json_array_elements(plan->'pegging')->>'opplan' as integer) opplan
        from demand
        where name = %s
        )
      select min(dmd.due), min(startdate), max(enddate)
      from dmd
      inner join operationplan
      on dmd.opplan = operationplan.id
      and type <> 'STCK'
      ''', (args[0]))
    x = cursor.fetchone()
    (due, start, end) = x
    if not due:
      # This demand is unplanned
      request.report_startdate = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
      request.report_enddate = request.report_startdate + timedelta(days=1)
      request.report_bucket = None
      request.report_bucketlist = []
      return
    if not start:
      start = due
    if not end:
      end = due

    # Adjust the horizon
    if due > end:
      end = due
    if due < start:
      start = due
    end += timedelta(days=1)
    start -= timedelta(days=1)
    request.report_startdate = start.replace(hour=0, minute=0, second=0, microsecond=0)
    request.report_enddate = end.replace(hour=0, minute=0, second=0, microsecond=0)
    request.report_bucket = None
    request.report_bucketlist = []

  @classmethod
  def query(reportclass, request, basequery):
    # Execute the query
    basesql, baseparams = basequery.query.get_compiler(basequery.db).as_sql(with_col_aliases=False)
    cursor = connections[request.database].cursor()

    # Get current date and horizon
    horizon = (request.report_enddate - request.report_startdate).total_seconds() / 10000
    try:
      current = datetime.strptime(
        Parameter.objects.using(request.database).get(name="currentdate").value,
        "%Y-%m-%d %H:%M:%S"
        )
    except:
      current = datetime.now()
      current = current.replace(microsecond=0)

    # Collect demand due date, all operationplans and loaded resources
    query = '''
      with pegging as (
        select
          min(rownum) as rownum, min(due) as due, opplan, min(lvl) as lvl, sum(quantity) as quantity
        from (select
          row_number() over () as rownum, opplan, due, lvl, quantity
        from (select
          due,
          cast(json_array_elements(plan->'pegging')->>'opplan' as integer) as opplan,
          cast(json_array_elements(plan->'pegging')->>'level' as integer) as lvl,
          cast(json_array_elements(plan->'pegging')->>'quantity' as numeric) as quantity
          from demand
          where name = %s
          ) d1
          )d2
        group by opplan
        )
      select
        pegging.due, operationplan.name, pegging.lvl, ops.pegged,
        pegging.rownum, operationplan.startdate, operationplan.enddate, operationplan.quantity,
        operationplan.status, operationplanresource.resource, operationplan.type,
        case when operationplan.operation_id is not null then 1 else 0 end as show
      from pegging
      inner join operationplan
        on operationplan.id = pegging.opplan
      inner join (
        select name,
          min(rownum) as rownum,
          sum(pegging.quantity) as pegged
        from pegging
        inner join operationplan
          on pegging.opplan = operationplan.id
        group by operationplan.name
        ) ops
      on operationplan.name = ops.name
      left outer join operationplanresource
        on pegging.opplan = operationplanresource.operationplan_id
      order by ops.rownum, pegging.rownum
      '''
    cursor.execute(query, baseparams)

    # Build the Python result
    prevrec = None
    parents = {}
    for rec in cursor.fetchall():
      if not prevrec or rec[1] != prevrec['operation']:
        # Return prev operation
        if prevrec:
          if prevrec['depth'] < rec[2]:
            prevrec['leaf'] = 'false'
          yield prevrec
        # New operation
        prevrec = {
          'current': str(current),
          'operation': rec[1],
          'type': rec[10],
          'showdrilldown': rec[11],
          'depth': rec[2],
          'quantity': str(rec[3]),
          'due': round((rec[0] - request.report_startdate).total_seconds() / horizon, 3),
          'current': round((current - request.report_startdate).total_seconds() / horizon, 3),
          'parent': rec[2] and parents[rec[2] - 1] or None,
          'leaf': 'true',
          'expanded': 'true',
          'resource': rec[9] and [rec[9], ] or [],
          'operationplans': [{
             'operation': rec[1],
             'quantity': str(rec[7]),
             'x': round((rec[5] - request.report_startdate).total_seconds() / horizon, 3),
             'w': round((rec[6] - rec[5]).total_seconds() / horizon, 3),
             'startdate': str(rec[5]),
             'enddate': str(rec[6]),
             'status': rec[8],
             'id': rec[4]
             }]
          }
        parents[rec[2]] = rec[1]
      elif rec[4] != prevrec['operationplans'][-1]['id']:
        # Extra operationplan for the operation
        prevrec['operationplans'].append({
          'operation': rec[1],
          'quantity': str(rec[7]),
          'x': round((rec[5] - request.report_startdate).total_seconds() / horizon, 3),
          'w': round((rec[6] - rec[5]).total_seconds() / horizon, 3),
          'startdate': str(rec[5]),
          'enddate': str(rec[6]),
          'locked': rec[8],
          'id': rec[4]
          })
      elif rec[9] and not rec[9] in prevrec['resource']:
        # Extra resource loaded by the operationplan
        prevrec['resource'].append(rec[9])
    if prevrec:
      yield prevrec
<commit_msg>demand plan report - extra robust check<commit_after>#
# Copyright (C) 2007-2013 by frePPLe bvba
#
# This library is free software; you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero
# General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from datetime import datetime, timedelta

from django.contrib.admin.utils import unquote
from django.db import connections
from django.http import Http404
from django.utils.translation import ugettext_lazy as _

from freppledb.input.models import Demand
from freppledb.common.report import GridReport, GridFieldText, GridFieldNumber
from freppledb.common.models import Parameter

class ReportByDemand(GridReport):
  '''
  A list report to show peggings.
  '''
  template = 'output/pegging.html'
  title = _("Demand plan")
  filterable = False
  frozenColumns = 0
  editable = False
  default_sort = None
  hasTimeBuckets = True
  multiselect = False
  heightmargin = 87
  help_url = 'user-guide/user-interface/plan-analysis/demand-gantt-report.html'
  rows = (
    GridFieldText('depth', title=_('depth'), editable=False, sortable=False),
    GridFieldText('operation', title=_('operation'), editable=False, sortable=False, key=True, formatter='detail', extra='"role":"input/operation"'),
    GridFieldText('type', title=_('type'), editable=False, sortable=False, width=100),
    #GridFieldText('buffer', title=_('buffer'), formatter='buffer', editable=False, sortable=False),
    #GridFieldText('item', title=_('item'), formatter='item', editable=False, sortable=False),
    GridFieldText('resource', title=_('resource'), editable=False, sortable=False, extra='formatter:reslistfmt'),
    GridFieldNumber('quantity', title=_('quantity'), editable=False, sortable=False),
    GridFieldText('operationplans', width=1000, extra='formatter:ganttcell', editable=False, sortable=False),
    GridFieldText('parent', editable=False, sortable=False, hidden=True),
    GridFieldText('leaf', editable=False, sortable=False, hidden=True),
    GridFieldText('expanded', editable=False, sortable=False, hidden=True),
    GridFieldText('current', editable=False, sortable=False, hidden=True),
    GridFieldText('due', editable=False, sortable=False, hidden=True),
    GridFieldText('showdrilldown', editable=False, sortable=False, hidden=True),
    )

  @ classmethod
  def basequeryset(reportclass, request, args, kwargs):
    return Demand.objects.filter(name__exact=unquote(args[0])).values('name')

  @classmethod
  def getBuckets(reportclass, request, *args, **kwargs):
    # Get the earliest and latest operationplan, and the demand due date
    cursor = connections[request.database].cursor()
    cursor.execute('''
      with dmd as (
        select
          due,
          cast(json_array_elements(plan->'pegging')->>'opplan' as integer) opplan
        from demand
        where name = %s
        )
      select min(dmd.due), min(startdate), max(enddate)
      from dmd
      inner join operationplan
      on dmd.opplan = operationplan.id
      and type <> 'STCK'
      ''', (args[0]))
    x = cursor.fetchone()
    (due, start, end) = x
    if not due:
      # This demand is unplanned
      request.report_startdate = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
      request.report_enddate = request.report_startdate + timedelta(days=1)
      request.report_bucket = None
      request.report_bucketlist = []
      return
    if not start:
      start = due
    if not end:
      end = due

    # Adjust the horizon
    if due > end:
      end = due
    if due < start:
      start = due
    end += timedelta(days=1)
    start -= timedelta(days=1)
    request.report_startdate = start.replace(hour=0, minute=0, second=0, microsecond=0)
    request.report_enddate = end.replace(hour=0, minute=0, second=0, microsecond=0)
    request.report_bucket = None
    request.report_bucketlist = []

  @classmethod
  def query(reportclass, request, basequery):
    # Execute the query
    basesql, baseparams = basequery.query.get_compiler(basequery.db).as_sql(with_col_aliases=False)
    cursor = connections[request.database].cursor()

    # Get current date and horizon
    horizon = (request.report_enddate - request.report_startdate).total_seconds() / 10000
    try:
      current = datetime.strptime(
        Parameter.objects.using(request.database).get(name="currentdate").value,
        "%Y-%m-%d %H:%M:%S"
        )
    except:
      current = datetime.now()
      current = current.replace(microsecond=0)

    # Collect demand due date, all operationplans and loaded resources
    query = '''
      with pegging as (
        select
          min(rownum) as rownum, min(due) as due, opplan, min(lvl) as lvl, sum(quantity) as quantity
        from (select
          row_number() over () as rownum, opplan, due, lvl, quantity
        from (select
          due,
          cast(json_array_elements(plan->'pegging')->>'opplan' as integer) as opplan,
          cast(json_array_elements(plan->'pegging')->>'level' as integer) as lvl,
          cast(json_array_elements(plan->'pegging')->>'quantity' as numeric) as quantity
          from demand
          where name = %s
          ) d1
          )d2
        group by opplan
        )
      select
        pegging.due, operationplan.name, pegging.lvl, ops.pegged,
        pegging.rownum, operationplan.startdate, operationplan.enddate, operationplan.quantity,
        operationplan.status, operationplanresource.resource, operationplan.type,
        case when operationplan.operation_id is not null then 1 else 0 end as show
      from pegging
      inner join operationplan
        on operationplan.id = pegging.opplan
      inner join (
        select name,
          min(rownum) as rownum,
          sum(pegging.quantity) as pegged
        from pegging
        inner join operationplan
          on pegging.opplan = operationplan.id
        group by operationplan.name
        ) ops
      on operationplan.name = ops.name
      left outer join operationplanresource
        on pegging.opplan = operationplanresource.operationplan_id
      order by ops.rownum, pegging.rownum
      '''
    cursor.execute(query, baseparams)

    # Build the Python result
    prevrec = None
    parents = {}
    for rec in cursor.fetchall():
      if not prevrec or rec[1] != prevrec['operation']:
        # Return prev operation
        if prevrec:
          if prevrec['depth'] < rec[2]:
            prevrec['leaf'] = 'false'
          yield prevrec
        # New operation
        prevrec = {
          'current': str(current),
          'operation': rec[1],
          'type': rec[10],
          'showdrilldown': rec[11],
          'depth': rec[2],
          'quantity': str(rec[3]),
          'due': round((rec[0] - request.report_startdate).total_seconds() / horizon, 3),
          'current': round((current - request.report_startdate).total_seconds() / horizon, 3),
          'parent': parents[rec[2] - 1] if rec[2] and rec[2] >= 1 and rec[2] < len(parents) else None,
          'leaf': 'true',
          'expanded': 'true',
          'resource': rec[9] and [rec[9], ] or [],
          'operationplans': [{
             'operation': rec[1],
             'quantity': str(rec[7]),
             'x': round((rec[5] - request.report_startdate).total_seconds() / horizon, 3),
             'w': round((rec[6] - rec[5]).total_seconds() / horizon, 3),
             'startdate': str(rec[5]),
             'enddate': str(rec[6]),
             'status': rec[8],
             'id': rec[4]
             }]
          }
        parents[rec[2]] = rec[1]
      elif rec[4] != prevrec['operationplans'][-1]['id']:
        # Extra operationplan for the operation
        prevrec['operationplans'].append({
          'operation': rec[1],
          'quantity': str(rec[7]),
          'x': round((rec[5] - request.report_startdate).total_seconds() / horizon, 3),
          'w': round((rec[6] - rec[5]).total_seconds() / horizon, 3),
          'startdate': str(rec[5]),
          'enddate': str(rec[6]),
          'locked': rec[8],
          'id': rec[4]
          })
      elif rec[9] and not rec[9] in prevrec['resource']:
        # Extra resource loaded by the operationplan
        prevrec['resource'].append(rec[9])
    if prevrec:
      yield prevrec
<|endoftext|>


<commit_before>import bpy
import bmesh
import struct
import json
import os

def uv_from_vert_first(uv_layer, v):
    for l in v.link_loops:
        uv_data = l[uv_layer]
        return uv_data.uv
    return None

def has_material_index(vertex, index):
    has_index = False
    for face in vertex.link_faces:
        if face.material_index == index:
            has_index = True
            return has_index
    return has_index

def round_3d(v):
    return round(v[0], 6), round(v[1], 6), round(v[2], 6)

def round_2d(v):
    return round(v[0], 6), round(v[1], 6)

def mesh_path(blender_object):
    name = blender_object.data.name
    if len(blender_object.modifiers) > 0:
        name = blender_object.name
    for modifier in blender_object.modifiers:
        name += "_" + modifier.name.lower()

    library = os.path.splitext(bpy.path.basename(bpy.context.blend_data.filepath))[0] + '/'
    if blender_object.data.library:
        library, file_extension = os.path.splitext(bpy.path.basename(blender_object.data.library.filepath))
        library = library + '/'
    path = library + "meshes/" + name + ".mesh"
    return path.strip('/')

def write_mesh_file(blender_object, write_dir):
    try:
        mesh = blender_object.to_mesh(depsgraph=bpy.context.depsgraph,
                                      apply_modifiers=True)

    except Exception as exception:
        print("Error in object " + blender_object.name)
        raise exception

    filepath = write_dir + '/' + mesh_path(blender_object)

    #bm = bmesh.new()
    #bm.from_mesh(mesh)

    #bmesh.ops.triangulate(bm, faces=bm.faces, quad_method=2)

    positions = []
    normals = []
    tangents = []
    texture_uvs = []
    weights = []

    faces = []
    vertex_dict = {}
    vertex_count = 0

    if len(mesh.uv_layers) >= 1:
        # for i, f in enumerate(mesh.tessfaces):
        for i, f in enumerate(mesh.loop_triangles):
            temp_faces = []
            for j, v in enumerate(f.vertices):
                position = round_3d(mesh.vertices[v].co.to_tuple())
                if f.use_smooth:
                    normal = round_3d(mesh.vertices[v].normal)
                else:
                    normal = round_3d(f.normal.to_tuple())
                texture_uv = list(round_2d(mesh.tessface_uv_textures[0].data[i].uv[j][0:2]))
                texture_uv[1] = 1.0 - texture_uv[1]
                texture_uv = tuple(texture_uv) #TODO: Not nice
                weight = mesh.vertices[v].bevel_weight;

                key = mesh.vertices[v].index
                vertex_index = vertex_dict.get(key)

                if f.use_smooth:
                    if vertex_index is None:  # vertex not found
                        vertex_dict[key] = vertex_count
                        positions.append(position)
                        normals.append(normal)
                        texture_uvs.append(texture_uv)
                        temp_faces.append(vertex_count)
                        tangents.append((0.0, 0.0, 0.0))
                        weights.append(weight)
                        vertex_count += 1
                    else:
                        inx = vertex_dict[key]
                        temp_faces.append(inx)
                else:
                    positions.append(position)
                    normals.append(normal)
                    texture_uvs.append(texture_uv)
                    temp_faces.append(vertex_count)
                    tangents.append((0.0, 0.0, 0.0))
                    weights.append(weight)
                    vertex_count += 1

            if len(temp_faces) == 3:
                faces.append(temp_faces)
            else:
                faces.append([temp_faces[0], temp_faces[1], temp_faces[2]])
                faces.append([temp_faces[0], temp_faces[2], temp_faces[3]])

        indices = [val for sublist in faces for val in sublist]
    else:
        raise Exception(mesh.name + " must have one uv layer")
    #bm.free()
    #del bm

    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    mesh_file = open(filepath, 'bw')

    # Header
    mesh_file.write(struct.pack('i', len(positions)))
    mesh_file.write(struct.pack('i', len(indices)))

    # Body
    for v in zip(positions, normals, tangents, texture_uvs, weights):
        mesh_file.write(struct.pack('fff', *v[0]))
        mesh_file.write(struct.pack('fff', *v[1]))
        mesh_file.write(struct.pack('fff', *v[2]))
        mesh_file.write(struct.pack('ff', *v[3]))
        mesh_file.write(struct.pack('f', v[4]))

    for i in indices:
        mesh_file.write(struct.pack('I', i))

    mesh_file.close()
    print('Wrote: ' + filepath)

def write(write_dir, objects):
    print("Writing meshes.")
    objects = [o for o in objects if o.type == 'MESH']

    for blender_object in objects:
        write_mesh_file(blender_object, write_dir)
<commit_msg>Mesh fix.<commit_after>import bpy
import bmesh
import struct
import json
import os

def uv_from_vert_first(uv_layer, v):
    for l in v.link_loops:
        uv_data = l[uv_layer]
        return uv_data.uv
    return None

def has_material_index(vertex, index):
    has_index = False
    for face in vertex.link_faces:
        if face.material_index == index:
            has_index = True
            return has_index
    return has_index

def round_3d(v):
    return round(v[0], 6), round(v[1], 6), round(v[2], 6)

def round_2d(v):
    return round(v[0], 6), round(v[1], 6)

def mesh_path(blender_object):
    name = blender_object.data.name
    if len(blender_object.modifiers) > 0:
        name = blender_object.name
    for modifier in blender_object.modifiers:
        name += "_" + modifier.name.lower()

    library = os.path.splitext(bpy.path.basename(bpy.context.blend_data.filepath))[0] + '/'
    if blender_object.data.library:
        library, file_extension = os.path.splitext(bpy.path.basename(blender_object.data.library.filepath))
        library = library + '/'
    path = library + "meshes/" + name + ".mesh"
    return path.strip('/')

def write_mesh_file(blender_object, write_dir):
    try:
        mesh = blender_object.to_mesh(depsgraph=bpy.context.depsgraph,
                                      apply_modifiers=True)

    except Exception as exception:
        print("Error in object " + blender_object.name)
        raise exception

    filepath = write_dir + '/' + mesh_path(blender_object)

    #bm = bmesh.new()
    #bm.from_mesh(mesh)

    #bmesh.ops.triangulate(bm, faces=bm.faces, quad_method=2)

    positions = []
    normals = []
    tangents = []
    texture_uvs = []
    weights = []

    faces = []
    vertex_dict = {}
    vertex_count = 0

    print(len(mesh.uv_layers))
    if len(mesh.uv_layers) >= 1:
        # for i, f in enumerate(mesh.tessfaces):
        mesh.calc_loop_triangles()
        print(len(mesh.loop_triangles))
        for i, f in enumerate(mesh.loop_triangles):
            temp_faces = []
            for j, v in enumerate(f.vertices):
                print(len(f.vertices))
                position = round_3d(mesh.vertices[v].co.to_tuple())
                if f.use_smooth:
                    normal = round_3d(mesh.vertices[v].normal)
                else:
                    normal = round_3d(f.normal.to_tuple())
                texture_uv = list(round_2d(mesh.tessface_uv_textures[0].data[i].uv[j][0:2]))
                texture_uv[1] = 1.0 - texture_uv[1]
                texture_uv = tuple(texture_uv) #TODO: Not nice
                weight = mesh.vertices[v].bevel_weight;

                key = mesh.vertices[v].index
                vertex_index = vertex_dict.get(key)

                if f.use_smooth:
                    if vertex_index is None:  # vertex not found
                        vertex_dict[key] = vertex_count
                        positions.append(position)
                        normals.append(normal)
                        texture_uvs.append(texture_uv)
                        temp_faces.append(vertex_count)
                        tangents.append((0.0, 0.0, 0.0))
                        weights.append(weight)
                        vertex_count += 1
                    else:
                        inx = vertex_dict[key]
                        temp_faces.append(inx)
                else:
                    positions.append(position)
                    normals.append(normal)
                    texture_uvs.append(texture_uv)
                    temp_faces.append(vertex_count)
                    tangents.append((0.0, 0.0, 0.0))
                    weights.append(weight)
                    vertex_count += 1

            if len(temp_faces) == 3:
                faces.append(temp_faces)
            else:
                faces.append([temp_faces[0], temp_faces[1], temp_faces[2]])
                faces.append([temp_faces[0], temp_faces[2], temp_faces[3]])

        indices = [val for sublist in faces for val in sublist]
    else:
        raise Exception(mesh.name + " must have one uv layer")
    #bm.free()
    #del bm

    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    mesh_file = open(filepath, 'bw')

    # Header
    mesh_file.write(struct.pack('i', len(positions)))
    mesh_file.write(struct.pack('i', len(indices)))

    # Body
    for v in zip(positions, normals, tangents, texture_uvs, weights):
        mesh_file.write(struct.pack('fff', *v[0]))
        mesh_file.write(struct.pack('fff', *v[1]))
        mesh_file.write(struct.pack('fff', *v[2]))
        mesh_file.write(struct.pack('ff', *v[3]))
        mesh_file.write(struct.pack('f', v[4]))

    for i in indices:
        mesh_file.write(struct.pack('I', i))

    mesh_file.close()
    print('Wrote: ' + filepath)

def write(write_dir, objects):
    print("Writing meshes.")
    objects = [o for o in objects if o.type == 'MESH']

    for blender_object in objects:
        write_mesh_file(blender_object, write_dir)
<|endoftext|>


<commit_before>import numpy as np

from seisflows.tools import unix
from seisflows.tools.array import loadnpy, savenpy
from seisflows.tools.code import divides, exists, glob, irange, join
from seisflows.tools.config import ConfigObj, ParameterObj

OBJ = ConfigObj('SeisflowsObjects')
PAR = ParameterObj('SeisflowsParameters')
PATH = ParameterObj('SeisflowsPaths')

import system
import solver
import optimize
import preprocess
import postprocess

class inversion(object):
    """ Seismic inversion base class.

      Computes iterative model updates using any of a variety of possible
      settings, and provides a base class on top of which custom inversion
      strategies can be implemented.

      To allow overriding, separate methods are provided for each step in the
      inversion, including 'evaluate_function', 'evaluate_gradient',
      'compute_direction', and 'line_search' as well as generic 'initialize'
      and 'finalize' methods.

      Calls to forward and adjoint solvers are abstracted through the 'solver'
      interface so that various forward modeling packages can be used
      interchangeably.

      Commands for launching serial or parallel jobs on a workstation or cluster
      are abstracted through the 'system' interface so that the code can be
      easily ported.

      For assistance using this package, please browse comments or email
      rmodrak -at- princeton -dot- edu
    """

    def check(self):
        """ Checks parameters, paths, and dependencies
        """

        # check parameters
        if 'BEGIN' not in PAR:
            raise Exception

        if 'END' not in PAR:
            raise Exception

        if 'VERBOSE' not in PAR:
            setattr(PAR, 'VERBOSE', 1)

        # check paths
        if 'GLOBAL' not in PATH:
            raise Exception

        if 'LOCAL' not in PATH:
            setattr(PATH, 'LOCAL', None)

        if 'FUNC' not in PATH:
            setattr(PATH, 'FUNC', join(PATH.GLOBAL, 'func'))

        if 'GRAD' not in PATH:
            setattr(PATH, 'GRAD', join(PATH.GLOBAL, 'grad'))

        if 'HESS' not in PATH:
            setattr(PATH, 'HESS', join(PATH.GLOBAL, 'hess'))

        if 'OPTIMIZE' not in PATH:
            setattr(PATH, 'OPTIMIZE', join(PATH.GLOBAL, 'optimize'))

        # check input settings
        if 'DATA' not in PATH:
            setattr(PATH, 'DATA', None)

        if not exists(PATH.DATA):
            assert 'MODEL_TRUE' in PATH

        if 'MODEL_INIT' not in PATH:
            raise Exception

        # check output settings
        if 'OUTPUT' not in PATH:
            raise Exception

        if 'SAVEMODEL' not in PAR:
            setattr(PAR, 'SAVEMODEL', 1)

        if 'SAVEGRADIENT' not in PAR:
            setattr(PAR, 'SAVEGRADIENT', 0)

        if 'SAVEKERNELS' not in PAR:
            setattr(PAR, 'SAVEKERNELS', 0)

        if 'SAVETRACES' not in PAR:
            setattr(PAR, 'SAVETRACES', 0)

        if 'SAVERESIDUALS' not in PAR:
            setattr(PAR, 'SAVERESIDUALS', 0)

    def main(self):
        """ Carries out seismic inversion
        """
        self.setup()

        for self.iter in irange(PAR.BEGIN, PAR.END):
            optimize.iter = self.iter

            print "Starting iteration", self.iter
            self.initialize()

            print "Computing search direction"
            self.compute_direction()

            print "Computing step length"
            self.line_search()

            self.finalize()
            print ''

            if self.isdone:
                return

    def setup(self):
        """ Lays groundwork for inversion
        """
        if PAR.BEGIN == 1:
            # prepare directory structure
            unix.rm(PATH.GLOBAL)
            unix.mkdir(PATH.GLOBAL)

            preprocess.setup()
            postprocess.setup()
            optimize.setup()

            system.run('solver', 'setup',
                       hosts='all')

        else:
            # prepare optimization and solver directories
            optimize.setup()

            if PATH.LOCAL:
                system.run('solver', 'setup',
                           hosts='all')

    def initialize(self):
        """ Prepares for next model update iteration
        """
        isready = self.solver_status()
        if not isready:
            print 'Generating synthetics'

            self.prepare_model(path=PATH.GRAD, suffix='new')

            # forward simulation
            system.run('solver', 'eval_func',
                       hosts='all',
                       path=PATH.GRAD)

            self.sum_residuals(path=PATH.GRAD, suffix='new')

    def compute_direction(self):
        """ Computes search direction
        """
        self.evaluate_gradient()
        optimize.compute_direction()

    def line_search(self):
        """ Conducts line search in given search direction
        """
        optimize.initialize_search()

        for optimize.step in irange(1, PAR.SRCHMAX):
            isdone = self.search_status()

            if isdone == 1:
                optimize.finalize_search()
                break
            elif isdone == 0:
                optimize.compute_step()
                continue
            elif isdone == -1:
                self.isdone = -1
                print ' line search failed'

    def search_status(self):
        """ Checks line search status
        """
        if PAR.VERBOSE:
            print " trial step", optimize.step
        self.evaluate_function()
        isdone, isbest = optimize.search_status()

        if not PATH.LOCAL:
            if isbest and isdone:
                unix.mv(PATH.SOLVER, PATH.SOLVER + '_best')
            elif isbest:
                unix.rm(PATH.SOLVER + '_best')
                unix.cp(PATH.SOLVER, PATH.SOLVER + '_best')

        return isdone

    def evaluate_function(self):
        """ Calls forward solver and writes misfit
        """
        self.prepare_model(path=PATH.FUNC, suffix='try')

        # forward simulation
        system.run('solver', 'eval_func',
                   hosts='all',
                   path=PATH.FUNC)

        self.sum_residuals(path=PATH.FUNC, suffix='try')

    def evaluate_gradient(self):
        """ Calls adjoint solver and runs process_kernels
        """
        # adjoint simulation
        system.run('solver', 'eval_grad',
                   hosts='all',
                   path=PATH.GRAD,
                   export_traces=divides(self.iter, PAR.SAVETRACES))

        postprocess.process_kernels(
            path=PATH.GRAD)

    def finalize(self):
        """ Saves results from most recent model update iteration
        """
        if divides(self.iter, PAR.SAVEMODEL):
            self.save_model()

        if divides(self.iter, PAR.SAVEGRADIENT):
            self.save_gradient()

        if divides(self.iter, PAR.SAVEKERNELS):
            self.save_kernels()

        if divides(self.iter, PAR.SAVETRACES):
            self.save_traces()

        if divides(self.iter, PAR.SAVERESIDUALS):
            self.save_residuals()

        # clean up directories for next iteration
        if not PATH.LOCAL:
            unix.rm(PATH.GRAD)
            unix.mv(PATH.FUNC, PATH.GRAD)
            unix.mkdir(PATH.FUNC)

            unix.rm(PATH.SOLVER)
            unix.mv(PATH.SOLVER + '_best', PATH.SOLVER)

        else:
            unix.rm(PATH.GRAD)
            unix.rm(PATH.FUNC)
            unix.mkdir(PATH.GRAD)
            unix.mkdir(PATH.FUNC)

        self.isdone = False

    # -- utility functions

    def prepare_model(self, path='', suffix=''):
        """ Writes model in format used by solver
        """
        unix.mkdir(path)
        src = PATH.OPTIMIZE +'/'+ 'm_' + suffix
        dst = path +'/'+ 'model'
        parts = solver.split(loadnpy(src))
        solver.save(dst, parts)

    def sum_residuals(self, path='', suffix=''):
        """ Sums residuals to obtain misfit function value
        """
        src = path +'/'+ 'residuals'
        dst = PATH.OPTIMIZE +'/'+ 'f_' + suffix
        residuals = []
        for file in unix.ls(src):
            fromfile = np.loadtxt(src +'/'+ file)
            residuals.append(fromfile**2.)
        np.savetxt(dst, [np.sum(residuals)])

    def solver_status(self):
        """ Decides if solver prerequisites are in place
        """
        if self.iter <= 1:
            isready = False
        elif PATH.LOCAL:
            isready = False
        else:
            isready = True
        return isready

    def save_model(self):
        src = PATH.OPTIMIZE +'/'+ 'm_new'
        dst = join(PATH.OUTPUT, 'model_%04d' % self.iter)
        solver.save(dst, solver.split(loadnpy(src)))

    def save_gradient(self):
        src = glob(join(PATH.POSTPROCESS, 'grad*'))
        dst = join(PATH.OUTPUT, 'gradient_%04d' % self.iter)
        unix.mv(src, dst)

    def save_kernels(self):
        src = join(PATH.GRAD, 'kernels')
        dst = join(PATH.OUTPUT, 'kernels_%04d' % self.iter)
        unix.mkdir(dst)
        unix.mv(src, dst)

    def save_traces(self):
        src = join(PATH.GRAD, 'traces')
        dst = join(PATH.OUTPUT, 'traces_%04d' % self.iter)
        unix.mv(src, dst)

    def save_residuals(self):
        src = join(PATH.GRAD, 'residuals')
        dst = join(PATH.OUTPUT, 'residuals_%04d' % self.iter)
        unix.mv(src, dst)
<commit_msg>tweaked inversion paths<commit_after>import numpy as np

from seisflows.tools import unix
from seisflows.tools.array import loadnpy, savenpy
from seisflows.tools.code import divides, exists, glob, irange, join
from seisflows.tools.config import ConfigObj, ParameterObj

OBJ = ConfigObj('SeisflowsObjects')
PAR = ParameterObj('SeisflowsParameters')
PATH = ParameterObj('SeisflowsPaths')

import system
import solver
import optimize
import preprocess
import postprocess

class inversion(object):
    """ Seismic inversion base class.

      Computes iterative model updates using any of a variety of possible
      settings, and provides a base class on top of which custom inversion
      strategies can be implemented.

      To allow overriding, separate methods are provided for each step in the
      inversion, including 'evaluate_function', 'evaluate_gradient',
      'compute_direction', and 'line_search' as well as generic 'initialize'
      and 'finalize' methods.

      Calls to forward and adjoint solvers are abstracted through the 'solver'
      interface so that various forward modeling packages can be used
      interchangeably.

      Commands for launching serial or parallel jobs on a workstation or cluster
      are abstracted through the 'system' interface so that the code can be
      easily ported.

      For assistance using this package, please browse comments or email
      rmodrak -at- princeton -dot- edu
    """

    def check(self):
        """ Checks parameters, paths, and dependencies
        """

        # check parameters
        if 'BEGIN' not in PAR:
            raise Exception

        if 'END' not in PAR:
            raise Exception

        if 'VERBOSE' not in PAR:
            setattr(PAR, 'VERBOSE', 1)

        # check paths
        if 'GLOBAL' not in PATH:
            raise Exception

        if 'LOCAL' not in PATH:
            setattr(PATH, 'LOCAL', None)

        if 'FUNC' not in PATH:
            setattr(PATH, 'FUNC', join(PATH.GLOBAL, 'evalfunc'))

        if 'GRAD' not in PATH:
            setattr(PATH, 'GRAD', join(PATH.GLOBAL, 'evalgrad'))

        if 'HESS' not in PATH:
            setattr(PATH, 'HESS', join(PATH.GLOBAL, 'evalhess'))

        if 'OPTIMIZE' not in PATH:
            setattr(PATH, 'OPTIMIZE', join(PATH.GLOBAL, 'optimize'))

        # check input settings
        if 'DATA' not in PATH:
            setattr(PATH, 'DATA', None)

        if not exists(PATH.DATA):
            assert 'MODEL_TRUE' in PATH

        if 'MODEL_INIT' not in PATH:
            raise Exception

        # check output settings
        if 'OUTPUT' not in PATH:
            raise Exception

        if 'SAVEMODEL' not in PAR:
            setattr(PAR, 'SAVEMODEL', 1)

        if 'SAVEGRADIENT' not in PAR:
            setattr(PAR, 'SAVEGRADIENT', 0)

        if 'SAVEKERNELS' not in PAR:
            setattr(PAR, 'SAVEKERNELS', 0)

        if 'SAVETRACES' not in PAR:
            setattr(PAR, 'SAVETRACES', 0)

        if 'SAVERESIDUALS' not in PAR:
            setattr(PAR, 'SAVERESIDUALS', 0)

    def main(self):
        """ Carries out seismic inversion
        """
        self.setup()

        for self.iter in irange(PAR.BEGIN, PAR.END):
            optimize.iter = self.iter

            print "Starting iteration", self.iter
            self.initialize()

            print "Computing search direction"
            self.compute_direction()

            print "Computing step length"
            self.line_search()

            self.finalize()
            print ''

            if self.isdone:
                return

    def setup(self):
        """ Lays groundwork for inversion
        """
        if PAR.BEGIN == 1:
            # prepare directory structure
            unix.rm(PATH.GLOBAL)
            unix.mkdir(PATH.GLOBAL)

            preprocess.setup()
            postprocess.setup()
            optimize.setup()

            system.run('solver', 'setup',
                       hosts='all')

        else:
            # prepare optimization and solver directories
            optimize.setup()

            if PATH.LOCAL:
                system.run('solver', 'setup',
                           hosts='all')

    def initialize(self):
        """ Prepares for next model update iteration
        """
        isready = self.solver_status()
        if not isready:
            print 'Generating synthetics'

            self.prepare_model(path=PATH.GRAD, suffix='new')

            # forward simulation
            system.run('solver', 'eval_func',
                       hosts='all',
                       path=PATH.GRAD)

            self.sum_residuals(path=PATH.GRAD, suffix='new')

    def compute_direction(self):
        """ Computes search direction
        """
        self.evaluate_gradient()
        optimize.compute_direction()

    def line_search(self):
        """ Conducts line search in given search direction
        """
        optimize.initialize_search()

        for optimize.step in irange(1, PAR.SRCHMAX):
            isdone = self.search_status()

            if isdone == 1:
                optimize.finalize_search()
                break
            elif isdone == 0:
                optimize.compute_step()
                continue
            elif isdone == -1:
                self.isdone = -1
                print ' line search failed'

    def search_status(self):
        """ Checks line search status
        """
        if PAR.VERBOSE:
            print " trial step", optimize.step
        self.evaluate_function()
        isdone, isbest = optimize.search_status()

        if not PATH.LOCAL:
            if isbest and isdone:
                unix.mv(PATH.SOLVER, PATH.SOLVER + '_best')
            elif isbest:
                unix.rm(PATH.SOLVER + '_best')
                unix.cp(PATH.SOLVER, PATH.SOLVER + '_best')

        return isdone

    def evaluate_function(self):
        """ Calls forward solver and writes misfit
        """
        self.prepare_model(path=PATH.FUNC, suffix='try')

        # forward simulation
        system.run('solver', 'eval_func',
                   hosts='all',
                   path=PATH.FUNC)

        self.sum_residuals(path=PATH.FUNC, suffix='try')

    def evaluate_gradient(self):
        """ Calls adjoint solver and runs process_kernels
        """
        # adjoint simulation
        system.run('solver', 'eval_grad',
                   hosts='all',
                   path=PATH.GRAD,
                   export_traces=divides(self.iter, PAR.SAVETRACES))

        postprocess.process_kernels(
            path=PATH.GRAD)

    def finalize(self):
        """ Saves results from most recent model update iteration
        """
        if divides(self.iter, PAR.SAVEMODEL):
            self.save_model()

        if divides(self.iter, PAR.SAVEGRADIENT):
            self.save_gradient()

        if divides(self.iter, PAR.SAVEKERNELS):
            self.save_kernels()

        if divides(self.iter, PAR.SAVETRACES):
            self.save_traces()

        if divides(self.iter, PAR.SAVERESIDUALS):
            self.save_residuals()

        # clean up directories for next iteration
        if not PATH.LOCAL:
            unix.rm(PATH.GRAD)
            unix.mv(PATH.FUNC, PATH.GRAD)
            unix.mkdir(PATH.FUNC)

            unix.rm(PATH.SOLVER)
            unix.mv(PATH.SOLVER + '_best', PATH.SOLVER)

        else:
            unix.rm(PATH.GRAD)
            unix.rm(PATH.FUNC)
            unix.mkdir(PATH.GRAD)
            unix.mkdir(PATH.FUNC)

        self.isdone = False

    # -- utility functions

    def prepare_model(self, path='', suffix=''):
        """ Writes model in format used by solver
        """
        unix.mkdir(path)
        src = PATH.OPTIMIZE +'/'+ 'm_' + suffix
        dst = path +'/'+ 'model'
        parts = solver.split(loadnpy(src))
        solver.save(dst, parts)

    def sum_residuals(self, path='', suffix=''):
        """ Sums residuals to obtain misfit function value
        """
        src = path +'/'+ 'residuals'
        dst = PATH.OPTIMIZE +'/'+ 'f_' + suffix
        residuals = []
        for file in unix.ls(src):
            fromfile = np.loadtxt(src +'/'+ file)
            residuals.append(fromfile**2.)
        np.savetxt(dst, [np.sum(residuals)])

    def solver_status(self):
        """ Decides if solver prerequisites are in place
        """
        if self.iter <= 1:
            isready = False
        elif PATH.LOCAL:
            isready = False
        else:
            isready = True
        return isready

    def save_model(self):
        src = PATH.OPTIMIZE +'/'+ 'm_new'
        dst = join(PATH.OUTPUT, 'model_%04d' % self.iter)
        solver.save(dst, solver.split(loadnpy(src)))

    def save_gradient(self):
        src = glob(join(PATH.POSTPROCESS, 'grad*'))
        dst = join(PATH.OUTPUT, 'gradient_%04d' % self.iter)
        unix.mv(src, dst)

    def save_kernels(self):
        src = join(PATH.GRAD, 'kernels')
        dst = join(PATH.OUTPUT, 'kernels_%04d' % self.iter)
        unix.mkdir(dst)
        unix.mv(src, dst)

    def save_traces(self):
        src = join(PATH.GRAD, 'traces')
        dst = join(PATH.OUTPUT, 'traces_%04d' % self.iter)
        unix.mv(src, dst)

    def save_residuals(self):
        src = join(PATH.GRAD, 'residuals')
        dst = join(PATH.OUTPUT, 'residuals_%04d' % self.iter)
        unix.mv(src, dst)
<|endoftext|>


<commit_before>from cereal import car
from common.conversions import Conversions as CV
from opendbc.can.can_define import CANDefine
from opendbc.can.parser import CANParser
from selfdrive.car.interfaces import CarStateBase
from selfdrive.car.mazda.values import DBC, LKAS_LIMITS, GEN1

class CarState(CarStateBase):
  def __init__(self, CP):
    super().__init__(CP)

    can_define = CANDefine(DBC[CP.carFingerprint]["pt"])
    self.shifter_values = can_define.dv["GEAR"]["GEAR"]

    self.crz_btns_counter = 0
    self.acc_active_last = False
    self.low_speed_alert = False
    self.lkas_allowed_speed = False
    self.lkas_disabled = False

  def update(self, cp, cp_cam):

    ret = car.CarState.new_message()
    ret.wheelSpeeds = self.get_wheel_speeds(
      cp.vl["WHEEL_SPEEDS"]["FL"],
      cp.vl["WHEEL_SPEEDS"]["FR"],
      cp.vl["WHEEL_SPEEDS"]["RL"],
      cp.vl["WHEEL_SPEEDS"]["RR"],
    )
    ret.vEgoRaw = (ret.wheelSpeeds.fl + ret.wheelSpeeds.fr + ret.wheelSpeeds.rl + ret.wheelSpeeds.rr) / 4.
    ret.vEgo, ret.aEgo = self.update_speed_kf(ret.vEgoRaw)

    # Match panda speed reading
    speed_kph = cp.vl["ENGINE_DATA"]["SPEED"]
    ret.standstill = speed_kph < .1

    can_gear = int(cp.vl["GEAR"]["GEAR"])
    ret.gearShifter = self.parse_gear_shifter(self.shifter_values.get(can_gear, None))

    ret.genericToggle = bool(cp.vl["BLINK_INFO"]["HIGH_BEAMS"])
    ret.leftBlindspot = cp.vl["BSM"]["LEFT_BS1"] == 1
    ret.rightBlindspot = cp.vl["BSM"]["RIGHT_BS1"] == 1
    ret.leftBlinker, ret.rightBlinker = self.update_blinker_from_lamp(40, cp.vl["BLINK_INFO"]["LEFT_BLINK"] == 1,
                                                                      cp.vl["BLINK_INFO"]["RIGHT_BLINK"] == 1)

    ret.steeringAngleDeg = cp.vl["STEER"]["STEER_ANGLE"]
    ret.steeringTorque = cp.vl["STEER_TORQUE"]["STEER_TORQUE_SENSOR"]
    ret.steeringPressed = abs(ret.steeringTorque) > LKAS_LIMITS.STEER_THRESHOLD

    ret.steeringTorqueEps = cp.vl["STEER_TORQUE"]["STEER_TORQUE_MOTOR"]
    ret.steeringRateDeg = cp.vl["STEER_RATE"]["STEER_ANGLE_RATE"]

    # TODO: this should be from 0 - 1.
    ret.brakePressed = cp.vl["PEDALS"]["BRAKE_ON"] == 1
    ret.brake = cp.vl["BRAKE"]["BRAKE_PRESSURE"]

    ret.seatbeltUnlatched = cp.vl["SEATBELT"]["DRIVER_SEATBELT"] == 0
    ret.doorOpen = any([cp.vl["DOORS"]["FL"], cp.vl["DOORS"]["FR"],
                        cp.vl["DOORS"]["BL"], cp.vl["DOORS"]["BR"]])

    # TODO: this should be from 0 - 1.
    ret.gas = cp.vl["ENGINE_DATA"]["PEDAL_GAS"]
    ret.gasPressed = ret.gas > 0

    # Either due to low speed or hands off
    lkas_blocked = cp.vl["STEER_RATE"]["LKAS_BLOCK"] == 1

    if self.CP.minSteerSpeed > 0:
      # LKAS is enabled at 52kph going up and disabled at 45kph going down
      # wait for LKAS_BLOCK signal to clear when going up since it lags behind the speed sometimes
      if speed_kph > LKAS_LIMITS.ENABLE_SPEED and not lkas_blocked:
        self.lkas_allowed_speed = True
      elif speed_kph < LKAS_LIMITS.DISABLE_SPEED:
        self.lkas_allowed_speed = False
    else:
      self.lkas_allowed_speed = True

    # TODO: the signal used for available seems to be the adaptive cruise signal, instead of the main on
    #       it should be used for carState.cruiseState.nonAdaptive instead
    ret.cruiseState.available = cp.vl["CRZ_CTRL"]["CRZ_AVAILABLE"] == 1
    ret.cruiseState.enabled = cp.vl["CRZ_CTRL"]["CRZ_ACTIVE"] == 1
    ret.cruiseState.speed = cp.vl["CRZ_EVENTS"]["CRZ_SPEED"] * CV.KPH_TO_MS

    if ret.cruiseState.enabled:
      if not self.lkas_allowed_speed and self.acc_active_last:
        self.low_speed_alert = True
      else:
        self.low_speed_alert = False

    # Check if LKAS is disabled due to lack of driver torque when all other states indicate
    # it should be enabled (steer lockout). Don't warn until we actually get lkas active
    # and lose it again, i.e, after initial lkas activation
    ret.steerFaultTemporary = self.lkas_allowed_speed and lkas_blocked

    self.acc_active_last = ret.cruiseState.enabled

    self.crz_btns_counter = cp.vl["CRZ_BTNS"]["CTR"]

    # camera signals
    self.lkas_disabled = cp_cam.vl["CAM_LANEINFO"]["LANE_LINES"] == 0
    self.cam_lkas = cp_cam.vl["CAM_LKAS"]
    self.cam_laneinfo = cp_cam.vl["CAM_LANEINFO"]
    ret.steerFaultPermanent = cp_cam.vl["CAM_LKAS"]["ERR_BIT_1"] == 1

    return ret

  @staticmethod
  def get_can_parser(CP):
    signals = [
      # sig_name, sig_address
      ("LEFT_BLINK", "BLINK_INFO"),
      ("RIGHT_BLINK", "BLINK_INFO"),
      ("HIGH_BEAMS", "BLINK_INFO"),
      ("STEER_ANGLE", "STEER"),
      ("STEER_ANGLE_RATE", "STEER_RATE"),
      ("STEER_TORQUE_SENSOR", "STEER_TORQUE"),
      ("STEER_TORQUE_MOTOR", "STEER_TORQUE"),
      ("FL", "WHEEL_SPEEDS"),
      ("FR", "WHEEL_SPEEDS"),
      ("RL", "WHEEL_SPEEDS"),
      ("RR", "WHEEL_SPEEDS"),
    ]

    checks = [
      # sig_address, frequency
      ("BLINK_INFO", 10),
      ("STEER", 67),
      ("STEER_RATE", 83),
      ("STEER_TORQUE", 83),
      ("WHEEL_SPEEDS", 100),
    ]

    if CP.carFingerprint in GEN1:
      signals += [
        ("LKAS_BLOCK", "STEER_RATE"),
        ("LKAS_TRACK_STATE", "STEER_RATE"),
        ("HANDS_OFF_5_SECONDS", "STEER_RATE"),
        ("CRZ_ACTIVE", "CRZ_CTRL"),
        ("CRZ_AVAILABLE", "CRZ_CTRL"),
        ("CRZ_SPEED", "CRZ_EVENTS"),
        ("STANDSTILL", "PEDALS"),
        ("BRAKE_ON", "PEDALS"),
        ("BRAKE_PRESSURE", "BRAKE"),
        ("GEAR", "GEAR"),
        ("DRIVER_SEATBELT", "SEATBELT"),
        ("FL", "DOORS"),
        ("FR", "DOORS"),
        ("BL", "DOORS"),
        ("BR", "DOORS"),
        ("PEDAL_GAS", "ENGINE_DATA"),
        ("SPEED", "ENGINE_DATA"),
        ("CTR", "CRZ_BTNS"),
        ("LEFT_BS1", "BSM"),
        ("RIGHT_BS1", "BSM"),
      ]

      checks += [
        ("ENGINE_DATA", 100),
        ("CRZ_CTRL", 50),
        ("CRZ_EVENTS", 50),
        ("CRZ_BTNS", 10),
        ("PEDALS", 50),
        ("BRAKE", 50),
        ("SEATBELT", 10),
        ("DOORS", 10),
        ("GEAR", 20),
        ("BSM", 10),
      ]

    return CANParser(DBC[CP.carFingerprint]["pt"], signals, checks, 0)

  @staticmethod
  def get_cam_can_parser(CP):
    signals = []
    checks = []

    if CP.carFingerprint in GEN1:
      signals += [
        # sig_name, sig_address
        ("LKAS_REQUEST", "CAM_LKAS"),
        ("CTR", "CAM_LKAS"),
        ("ERR_BIT_1", "CAM_LKAS"),
        ("LINE_NOT_VISIBLE", "CAM_LKAS"),
        ("BIT_1", "CAM_LKAS"),
        ("ERR_BIT_2", "CAM_LKAS"),
        ("STEERING_ANGLE", "CAM_LKAS"),
        ("ANGLE_ENABLED", "CAM_LKAS"),
        ("CHKSUM", "CAM_LKAS"),

        ("LINE_VISIBLE", "CAM_LANEINFO"),
        ("LINE_NOT_VISIBLE", "CAM_LANEINFO"),
        ("LANE_LINES", "CAM_LANEINFO"),
        ("BIT1", "CAM_LANEINFO"),
        ("BIT2", "CAM_LANEINFO"),
        ("BIT3", "CAM_LANEINFO"),
        ("NO_ERR_BIT", "CAM_LANEINFO"),
        ("S1", "CAM_LANEINFO"),
        ("S1_HBEAM", "CAM_LANEINFO"),
      ]

      checks += [
        # sig_address, frequency
        ("CAM_LANEINFO", 2),
        ("CAM_LKAS", 16),
      ]

    return CANParser(DBC[CP.carFingerprint]["pt"], signals, checks, 2)
<commit_msg>Mazda: log standstill bit (#24923)<commit_after>from cereal import car
from common.conversions import Conversions as CV
from opendbc.can.can_define import CANDefine
from opendbc.can.parser import CANParser
from selfdrive.car.interfaces import CarStateBase
from selfdrive.car.mazda.values import DBC, LKAS_LIMITS, GEN1

class CarState(CarStateBase):
  def __init__(self, CP):
    super().__init__(CP)

    can_define = CANDefine(DBC[CP.carFingerprint]["pt"])
    self.shifter_values = can_define.dv["GEAR"]["GEAR"]

    self.crz_btns_counter = 0
    self.acc_active_last = False
    self.low_speed_alert = False
    self.lkas_allowed_speed = False
    self.lkas_disabled = False

  def update(self, cp, cp_cam):

    ret = car.CarState.new_message()
    ret.wheelSpeeds = self.get_wheel_speeds(
      cp.vl["WHEEL_SPEEDS"]["FL"],
      cp.vl["WHEEL_SPEEDS"]["FR"],
      cp.vl["WHEEL_SPEEDS"]["RL"],
      cp.vl["WHEEL_SPEEDS"]["RR"],
    )
    ret.vEgoRaw = (ret.wheelSpeeds.fl + ret.wheelSpeeds.fr + ret.wheelSpeeds.rl + ret.wheelSpeeds.rr) / 4.
    ret.vEgo, ret.aEgo = self.update_speed_kf(ret.vEgoRaw)

    # Match panda speed reading
    speed_kph = cp.vl["ENGINE_DATA"]["SPEED"]
    ret.standstill = speed_kph < .1

    can_gear = int(cp.vl["GEAR"]["GEAR"])
    ret.gearShifter = self.parse_gear_shifter(self.shifter_values.get(can_gear, None))

    ret.genericToggle = bool(cp.vl["BLINK_INFO"]["HIGH_BEAMS"])
    ret.leftBlindspot = cp.vl["BSM"]["LEFT_BS1"] == 1
    ret.rightBlindspot = cp.vl["BSM"]["RIGHT_BS1"] == 1
    ret.leftBlinker, ret.rightBlinker = self.update_blinker_from_lamp(40, cp.vl["BLINK_INFO"]["LEFT_BLINK"] == 1,
                                                                      cp.vl["BLINK_INFO"]["RIGHT_BLINK"] == 1)

    ret.steeringAngleDeg = cp.vl["STEER"]["STEER_ANGLE"]
    ret.steeringTorque = cp.vl["STEER_TORQUE"]["STEER_TORQUE_SENSOR"]
    ret.steeringPressed = abs(ret.steeringTorque) > LKAS_LIMITS.STEER_THRESHOLD

    ret.steeringTorqueEps = cp.vl["STEER_TORQUE"]["STEER_TORQUE_MOTOR"]
    ret.steeringRateDeg = cp.vl["STEER_RATE"]["STEER_ANGLE_RATE"]

    # TODO: this should be from 0 - 1.
    ret.brakePressed = cp.vl["PEDALS"]["BRAKE_ON"] == 1
    ret.brake = cp.vl["BRAKE"]["BRAKE_PRESSURE"]

    ret.seatbeltUnlatched = cp.vl["SEATBELT"]["DRIVER_SEATBELT"] == 0
    ret.doorOpen = any([cp.vl["DOORS"]["FL"], cp.vl["DOORS"]["FR"],
                        cp.vl["DOORS"]["BL"], cp.vl["DOORS"]["BR"]])

    # TODO: this should be from 0 - 1.
    ret.gas = cp.vl["ENGINE_DATA"]["PEDAL_GAS"]
    ret.gasPressed = ret.gas > 0

    # Either due to low speed or hands off
    lkas_blocked = cp.vl["STEER_RATE"]["LKAS_BLOCK"] == 1

    if self.CP.minSteerSpeed > 0:
      # LKAS is enabled at 52kph going up and disabled at 45kph going down
      # wait for LKAS_BLOCK signal to clear when going up since it lags behind the speed sometimes
      if speed_kph > LKAS_LIMITS.ENABLE_SPEED and not lkas_blocked:
        self.lkas_allowed_speed = True
      elif speed_kph < LKAS_LIMITS.DISABLE_SPEED:
        self.lkas_allowed_speed = False
    else:
      self.lkas_allowed_speed = True

    # TODO: the signal used for available seems to be the adaptive cruise signal, instead of the main on
    #       it should be used for carState.cruiseState.nonAdaptive instead
    ret.cruiseState.available = cp.vl["CRZ_CTRL"]["CRZ_AVAILABLE"] == 1
    ret.cruiseState.enabled = cp.vl["CRZ_CTRL"]["CRZ_ACTIVE"] == 1
    ret.cruiseState.standstill = cp.vl["PEDALS"]["STANDSTILL"] == 1
    ret.cruiseState.speed = cp.vl["CRZ_EVENTS"]["CRZ_SPEED"] * CV.KPH_TO_MS

    if ret.cruiseState.enabled:
      if not self.lkas_allowed_speed and self.acc_active_last:
        self.low_speed_alert = True
      else:
        self.low_speed_alert = False

    # Check if LKAS is disabled due to lack of driver torque when all other states indicate
    # it should be enabled (steer lockout). Don't warn until we actually get lkas active
    # and lose it again, i.e, after initial lkas activation
    ret.steerFaultTemporary = self.lkas_allowed_speed and lkas_blocked

    self.acc_active_last = ret.cruiseState.enabled

    self.crz_btns_counter = cp.vl["CRZ_BTNS"]["CTR"]

    # camera signals
    self.lkas_disabled = cp_cam.vl["CAM_LANEINFO"]["LANE_LINES"] == 0
    self.cam_lkas = cp_cam.vl["CAM_LKAS"]
    self.cam_laneinfo = cp_cam.vl["CAM_LANEINFO"]
    ret.steerFaultPermanent = cp_cam.vl["CAM_LKAS"]["ERR_BIT_1"] == 1

    return ret

  @staticmethod
  def get_can_parser(CP):
    signals = [
      # sig_name, sig_address
      ("LEFT_BLINK", "BLINK_INFO"),
      ("RIGHT_BLINK", "BLINK_INFO"),
      ("HIGH_BEAMS", "BLINK_INFO"),
      ("STEER_ANGLE", "STEER"),
      ("STEER_ANGLE_RATE", "STEER_RATE"),
      ("STEER_TORQUE_SENSOR", "STEER_TORQUE"),
      ("STEER_TORQUE_MOTOR", "STEER_TORQUE"),
      ("FL", "WHEEL_SPEEDS"),
      ("FR", "WHEEL_SPEEDS"),
      ("RL", "WHEEL_SPEEDS"),
      ("RR", "WHEEL_SPEEDS"),
    ]

    checks = [
      # sig_address, frequency
      ("BLINK_INFO", 10),
      ("STEER", 67),
      ("STEER_RATE", 83),
      ("STEER_TORQUE", 83),
      ("WHEEL_SPEEDS", 100),
    ]

    if CP.carFingerprint in GEN1:
      signals += [
        ("LKAS_BLOCK", "STEER_RATE"),
        ("LKAS_TRACK_STATE", "STEER_RATE"),
        ("HANDS_OFF_5_SECONDS", "STEER_RATE"),
        ("CRZ_ACTIVE", "CRZ_CTRL"),
        ("CRZ_AVAILABLE", "CRZ_CTRL"),
        ("CRZ_SPEED", "CRZ_EVENTS"),
        ("STANDSTILL", "PEDALS"),
        ("BRAKE_ON", "PEDALS"),
        ("BRAKE_PRESSURE", "BRAKE"),
        ("GEAR", "GEAR"),
        ("DRIVER_SEATBELT", "SEATBELT"),
        ("FL", "DOORS"),
        ("FR", "DOORS"),
        ("BL", "DOORS"),
        ("BR", "DOORS"),
        ("PEDAL_GAS", "ENGINE_DATA"),
        ("SPEED", "ENGINE_DATA"),
        ("CTR", "CRZ_BTNS"),
        ("LEFT_BS1", "BSM"),
        ("RIGHT_BS1", "BSM"),
      ]

      checks += [
        ("ENGINE_DATA", 100),
        ("CRZ_CTRL", 50),
        ("CRZ_EVENTS", 50),
        ("CRZ_BTNS", 10),
        ("PEDALS", 50),
        ("BRAKE", 50),
        ("SEATBELT", 10),
        ("DOORS", 10),
        ("GEAR", 20),
        ("BSM", 10),
      ]

    return CANParser(DBC[CP.carFingerprint]["pt"], signals, checks, 0)

  @staticmethod
  def get_cam_can_parser(CP):
    signals = []
    checks = []

    if CP.carFingerprint in GEN1:
      signals += [
        # sig_name, sig_address
        ("LKAS_REQUEST", "CAM_LKAS"),
        ("CTR", "CAM_LKAS"),
        ("ERR_BIT_1", "CAM_LKAS"),
        ("LINE_NOT_VISIBLE", "CAM_LKAS"),
        ("BIT_1", "CAM_LKAS"),
        ("ERR_BIT_2", "CAM_LKAS"),
        ("STEERING_ANGLE", "CAM_LKAS"),
        ("ANGLE_ENABLED", "CAM_LKAS"),
        ("CHKSUM", "CAM_LKAS"),

        ("LINE_VISIBLE", "CAM_LANEINFO"),
        ("LINE_NOT_VISIBLE", "CAM_LANEINFO"),
        ("LANE_LINES", "CAM_LANEINFO"),
        ("BIT1", "CAM_LANEINFO"),
        ("BIT2", "CAM_LANEINFO"),
        ("BIT3", "CAM_LANEINFO"),
        ("NO_ERR_BIT", "CAM_LANEINFO"),
        ("S1", "CAM_LANEINFO"),
        ("S1_HBEAM", "CAM_LANEINFO"),
      ]

      checks += [
        # sig_address, frequency
        ("CAM_LANEINFO", 2),
        ("CAM_LKAS", 16),
      ]

    return CANParser(DBC[CP.carFingerprint]["pt"], signals, checks, 2)
<|endoftext|>